Merge pull request #5 from aws-samples/MO-7967

Adding code to support the delete_graph functionality

Author: smsilb

disableDetective.py

@@ -42,23 +42,28 @@ def _master_account_type(val: str, pattern: str = r'[0-9]{12}'):
         return val
 
     # Setup command line arguments
-    parser = argparse.ArgumentParser(description=('Link AWS Accounts to central '
-                                                  'Detective Account.'))
+    parser = argparse.ArgumentParser(description=('Unlink AWS Accounts from a central '
+                                                  'Detective Account, or delete the entire Detective graph.'))
     parser.add_argument('--master_account', type=_master_account_type,
                         required=True,
                         help="AccountId for Central AWS Account.")
     parser.add_argument('--input_file', type=argparse.FileType('r'),
-                        required=True,
                         help=('Path to CSV file containing the list of '
-                              'account IDs and Email addresses.'))
+                              'account IDs and Email addresses. '
+                              'This does not need to be provided if you use the delete_graph flag.'))
     parser.add_argument('--assume_role', type=str, required=True,
                         help="Role Name to assume in each account.")
-    parser.add_argument('--delete_master', action='store_true', default=False,
-                        help="Delete the master Detective graph.")
+    parser.add_argument('--delete_graph', action='store_true',
+                        help=('Delete the master Detective graph. '
+                              'If not provided, you must provide an input file.'))
     parser.add_argument('--disabled_regions', type=str,
                         help=('Regions to disable Detective. If not specified, '
                               'all available regions disabled.'))
-    return parser.parse_args(args)
+    args = parser.parse_args(args)
+    if not args.delete_graph and not args.input_file:
+        raise parser.error("Either an input file or the delete_graph flag should be provided.")
+
+    return args
 
 
 def read_accounts_csv(input_file: typing.IO) -> typing.Dict:
@@ -74,6 +79,9 @@ def read_accounts_csv(input_file: typing.IO) -> typing.Dict:
     account_re = re.compile(r'[0-9]{12}')
     aws_account_dict = {}
 
+    if not input_file:
+        return aws_account_dict
+
     for acct in input_file.readlines():
         split_line = acct.strip().split(',')
 
@@ -255,6 +263,12 @@ def delete_members(d_client: botocore.client.BaseClient, graph_arn: str,
     args = setup_command_line()
     aws_account_dict = read_accounts_csv(args.input_file)
 
+    # making sure that we either have an account list to delete from graphs or the delete_graph flag is provided.
+    if len(list(aws_account_dict.keys())) == 0 and not args.delete_graph:
+        logging.error("The delete_graph flag was False while the provided account list is empty. "\
+            "Please check your inputs and re-run the script.")
+        exit(1)
+
     try:
         session = boto3.session.Session()
         detective_regions = get_regions(session, args.disabled_regions)
@@ -277,15 +291,16 @@ def delete_members(d_client: botocore.client.BaseClient, graph_arn: str,
             d_client = master_session.client('detective', region_name=region)
             graphs = get_graphs(d_client)
             if not graphs:
-                logging.info(f'Amazon Detective is NOT disabled in {region}')
-                continue
-            logging.info(f'Amazon Detective is disabled in region {region}')
+                logging.info(f'Amazon Detective has already been disabled in {region}')
+            else:
+                logging.info(f'Disabling Amazon Detective in region {region}')
 
             try:
-
                 for graph in graphs:
-                    delete_members(
-                        d_client, graph, aws_account_dict)
+                    if not args.delete_graph:
+                        delete_members(d_client, graph, aws_account_dict)
+                    else:
+                        d_client.delete_graph(graph)
             except NameError as e:
                 logging.error(f'account is not defined: {e}')
             except Exception as e:

tests/test_scripts.py

@@ -57,10 +57,23 @@ def test_setup_command_line_disableDetective():
     with pytest.raises(SystemExit):
     	disableDetective.setup_command_line(['--master_account', '12345', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv'])
 
-    # Non existent input file
+    # Non existent input file and no delete_graph flag
     with pytest.raises(SystemExit):
     	disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts1.csv'])
 
+    # Non existent input file with delete_graph flag provided
+    with pytest.raises(SystemExit):
+    	disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts1.csv', '--delete_graph'])
+    
+    # No input file provided and no delete_graph flag
+    with pytest.raises(SystemExit):
+        disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1'])
+    
+    # No input file provided but delete_graph is set
+    args = disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--delete_graph'])
+    assert args.input_file == None
+    assert args.delete_graph == True
+    
 ###
 # The purpose of this test is to make sure we read accounts and emails correctly from the input .csv file in enableDetective.py
 ###
@@ -76,13 +89,25 @@ def test_read_accounts_csv_enableDetective():
 # The purpose of this test is to make sure we read accounts and emails correctly from the input .csv file in disableDetective.py
 ###
 def test_read_accounts_csv_disableDetective():
+	# a test case where an input file is provided, although some of the lines are not correct
 	args = disableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv'])
 
 	accounts_dict = disableDetective.read_accounts_csv(args.input_file)
 
 	assert len(accounts_dict.keys()) == 6
 	assert accounts_dict == {"123456789012":"random@gmail.com", "000012345678":"email@gmail.com", "555555555555":"test5@gmail.com", "111111111111":"test1@gmail.com", "222222222222":"test2@gmail.com", "333333333333":"test3@gmail.com"}
 
+	# a test case where no input file is provided
+	args = disableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--delete_graph'])
+	assert args.input_file == None
+	accounts_dict = disableDetective.read_accounts_csv(args.input_file)
+	assert accounts_dict == {}
+
+	# a test case where an empty input file is provided, along with delete_graph. This is not an error, although clients should not run with this kind of input
+	args = disableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts2.csv', '--delete_graph'])
+	accounts_dict = disableDetective.read_accounts_csv(args.input_file)
+	assert accounts_dict == {}
+	assert args.delete_graph == True
 
 ###
 # The purpose of this test is to make sure we extract regions correctly in enableDetective.py