Support the option to add tags to newly created Detective graphs

Author: abiola-aws

enableDetective.py

@@ -42,7 +42,14 @@ def _master_account_type(val: str, pattern: str = r'[0-9]{12}'):
             raise argparse.ArgumentTypeError
         return val
 
-    # Setup command line arguments
+    class ParseCommaSeparatedKeyValuePairsAction(argparse.Action):
+        def __call__(self, parser, namespace, values, option_string=None):
+            setattr(namespace, self.dest, dict())
+            for kv_pairs in values.split(","):
+                key, value = kv_pairs.split('=', 1)
+                getattr(namespace, self.dest)[key] = value
+
+   # Setup command line arguments
     parser = argparse.ArgumentParser(description=('Link AWS Accounts to central '
                                                   'Detective Account.'))
     parser.add_argument('--master_account', type=_master_account_type,
@@ -57,6 +64,9 @@ def _master_account_type(val: str, pattern: str = r'[0-9]{12}'):
     parser.add_argument('--enabled_regions', type=str,
                         help=('Regions to enable Detective. If not specified, '
                               'all available regions enabled.'))
+    parser.add_argument('--tags',
+                        action=ParseCommaSeparatedKeyValuePairsAction,
+                        help="Tags to be added to any newly enabled Detective graphs.")
     return parser.parse_args(args)
 
 
@@ -286,15 +296,15 @@ def accept_invitations(role: str, accounts: typing.Set[str], graph: str, region:
     except Exception as e:
         logging.exception(f'error accepting invitation {e.args}')
 
-def enable_detective(d_client: botocore.client.BaseClient, region: str):
+def enable_detective(d_client: botocore.client.BaseClient, region: str, tags: dict = None):
     graphs = get_graphs(d_client)
 
     if not graphs:
         confirm = input('Should Amazon Detective be enabled in {}? Enter [Y/N]: '.format(region))
 
         if confirm == 'Y' or confirm == 'y':
-            logging.info(f'Enabling Amazon Detective in {region}')
-            graphs = [d_client.create_graph()['GraphArn']]
+            logging.info(f'Enabling Amazon Detective in {region}' + (f'with tags {tags}' if tags else ''))
+            graphs = [d_client.create_graph(Tags=tags)['GraphArn']]
         else:
             logging.info(f'Skipping {region}')
             return None
@@ -326,7 +336,7 @@ def enable_detective(d_client: botocore.client.BaseClient, region: str):
     for region in detective_regions:
         try:
             d_client = master_session.client('detective', region_name=region)
-            graphs = enable_detective(d_client, region)
+            graphs = enable_detective(d_client, region, args.tags)
 
             if graphs is None:
                 continue

tests/test_scripts.py

@@ -24,6 +24,14 @@ def test_setup_command_line_enableDetective():
 
     args = enableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--enabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv'])
     assert args.master_account == '000000000001'
+    assert args.tags == None
+
+    args = enableDetective.setup_command_line("--master_account 123456789012 --assume_role detectiveAdmin --input_file accounts.csv --tags TagKey1=TagValue1,TagKey2=TagValue2,TagKey3=TagValue3".split(" "))
+    assert args.tags == {
+        "TagKey1": "TagValue1",
+        "TagKey2": "TagValue2",
+        "TagKey3": "TagValue3",
+    }
 
     # Wrong master account
     # The internal function _master_account_type() should raise argparse.ArgumentTypeError, however this exception gets supressed by argparse, and SystemExit is raised instead.