Refactor Account Collection for EUC (#321)

Co-authored-by: chodono-aws <chodono@amazon.com>

Author: iakov-gan

data-collection/deploy/account-collector.yaml

@@ -22,6 +22,11 @@ Parameters:
     Type: String
     Description: "ARNs of KMS Keys for data buckets and/or Glue Catalog. Comma separated list, no spaces. Keep empty if data Buckets and Glue Catalog are not Encrypted with KMS. You can also set it to '*' to grant decrypt permission for all the keys."
     Default: ""
+  EUCAccountIDs:
+    Type: String
+    Description: "Comma-separated list of account IDs that user has identified to have WorkSpaces running, used to collect WorkSpaces metrics from. If left blank, metrics will be checked from all linked accounts in Organization."
+    Default: ""
+
 Outputs:
   LambdaFunctionName:
     Value: !Ref LambdaFunction
@@ -133,9 +138,8 @@ Resources:
             MANAGEMENT_ACCOUNT_IDS = os.environ.get('MANAGEMENT_ACCOUNT_IDS')
             BUCKET = os.environ.get('BUCKET_NAME')
             PREDEF_ACCOUNT_LIST_KEY = os.environ.get('PREDEF_ACCOUNT_LIST_KEY')
-            LINKED_ACCOUNT_LIST_KEY = os.environ.get('LINKED_ACCOUNT_LIST_KEY')
-            PAYER_ACCOUNT_LIST_KEY = os.environ.get('PAYER_ACCOUNT_LIST_KEY')
             EXCLUDED_ACCOUNT_LIST_KEY = os.environ.get('EXCLUDED_ACCOUNT_LIST_KEY')
+            EUC_ACCOUNTS = os.environ.get('EUC_ACCOUNT_IDS', '').strip()
             TMP_FILE = "/tmp/data.json"
 
             logger = logging.getLogger(__name__)
@@ -154,6 +158,7 @@ Resources:
 
                 functions = { # keep keys same as boto3 services
                     'linked': iterate_linked_accounts,
+                    'euc': partial(iterate_accounts_with_filter, EUC_ACCOUNTS),
                     'payers': partial(iterate_admins_accounts, None),
                     'organizations': partial(iterate_admins_accounts, 'organizations'),
                     'compute-optimizer': partial(iterate_admins_accounts, 'compute-optimizer'),
@@ -177,7 +182,7 @@ Resources:
                 if count == 0:
                     raise Exception('No accounts found. Check the log.') #pylint: disable=broad-exception-raised
 
-                key = LINKED_ACCOUNT_LIST_KEY if account_type == 'linked' else PAYER_ACCOUNT_LIST_KEY
+                key = f"account-list/{account_type}-account-list.json"
                 s3 = boto3.client('s3')
                 s3.upload_file(TMP_FILE, Bucket=BUCKET, Key=key)
 
@@ -196,7 +201,7 @@ Resources:
                         try:
                             account_id = ssm.get_parameter(Name=ssm_key)['Parameter']['Value']
                         except ssm.exceptions.ParameterNotFound:
-                            logger.warning(f'Not found ssm parameter {ssm_key}. Will use Management Account Id {payer_id}')
+                            logger.info(f'Not found ssm parameter {ssm_key}. Will use Management Account Id {payer_id}')
                     yield {"account": json.dumps({'account_id': account_id, 'account_name': '', 'payer_id': payer_id})}
 
             def iterate_linked_accounts():
@@ -230,6 +235,15 @@ Resources:
                 except Exception as exc: #pylint: disable=broad-exception-caught
                     logger.error( f'{type(exc).__name__}: When trying to build linked account list. {exc} ')
 
+            def iterate_accounts_with_filter(filter_accounts):
+                ''' same as iterate_linked_accounts but with additional filtering
+                filter_accounts: a comma separated list of accounts OR '*'
+                '''
+                for acc in iterate_linked_accounts():
+                    account_id = json.loads(acc['account'])['account_id']
+                    if account_id in filter_accounts.split(',') or filter_accounts.strip() == '*':
+                        yield acc
+
             def get_defined_list(bucket, key):
                 s3 = boto3.client("s3")
                 exts = [".json", ".csv"]
@@ -282,9 +296,8 @@ Resources:
           RESOURCE_PREFIX: !Ref ResourcePrefix
           BUCKET_NAME: !Ref DestinationBucket
           PREDEF_ACCOUNT_LIST_KEY: "account-list/account-list"
-          LINKED_ACCOUNT_LIST_KEY: "account-list/linked-account-list.json"
-          PAYER_ACCOUNT_LIST_KEY: "account-list/payer-account-list.json"
           EXCLUDED_ACCOUNT_LIST_KEY: "account-list/excluded-linked-account-list.csv"
+          EUC_ACCOUNT_IDS: !Ref EUCAccountIDs
     Metadata:
       cfn_nag:
         rules_to_suppress:

data-collection/deploy/deploy-data-collection.yaml

@@ -30,22 +30,18 @@ Metadata:
           - IncludeInventoryCollectorModule
           - IncludeOrgDataModule
           - IncludeRDSUtilizationModule
-          - IncludeEUCUtilizationModule
           - IncludeRightsizingModule
           - IncludeTAModule
           - IncludeTransitGatewayModule
           - IncludeAWSFeedsModule
           - IncludeLicenseManagerModule
           - IncludeQuickSightModule
           - IncludeServiceQuotasModule
-      - Label:
-          default: 'EUC Module Configuration'
-        Parameters:
           - IncludeEUCUtilizationModule
       - Label:
-          default: 'EUC Module Settings'
+          default: 'EUC (End User Compute) Module Configuration'
         Parameters:
-          - EUCAccountIDs        
+          - EUCAccountIDs
     ParameterLabels:
       DestinationBucket:
         default: 'Destination S3 bucket prefix'
@@ -224,11 +220,11 @@ Parameters:
     Type: String
     Description: Collects WorkSpaces CloudWatch metrics from your accounts
     AllowedValues: ['yes', 'no']
-    Default: 'no'  
+    Default: 'no'
   EUCAccountIDs:
     Type: String
-    Description: "Optional, If you enable EUC Utilization or Inventory module and you use Amazon WorkSpaces, please provide a comma-separated list of account IDs where WorkSpaces are deployed. If left blank, metrics will be collected from all linked accounts in the Organization."
-    Default: ""
+    Description: "Optional, If you enable EUC Utilization or Inventory module and you use Amazon WorkSpaces, please provide a comma-separated list of account IDs where WorkSpaces are deployed. Or you can set * to collect from all linked accounts in the Organization."
+    Default: "*"
   IncludeOrgDataModule:
     Type: String
     Description: Collects AWS Organizations data such as account Id, account name, organization parent and specified tags
@@ -1203,7 +1199,6 @@ Resources:
         StepFunctionTemplate: !FindInMap [StepFunctionCode, main-v3, TemplatePath]
         StepFunctionExecutionRoleARN: !GetAtt StepFunctionExecutionRole.Arn
         SchedulerExecutionRoleARN: !GetAtt SchedulerExecutionRole.Arn
-        EUCAccountIDs: !Ref EUCAccountIDs
         RegionsInScope:
           Fn::If:
             - RegionsInScopeIsEmpty
@@ -1396,6 +1391,7 @@ Resources:
         DestinationBucket: !Ref S3Bucket
         DestinationBucketARN: !GetAtt S3Bucket.Arn
         DataBucketsKmsKeysArns: !Ref DataBucketsKmsKeysArns
+        EUCAccountIDs: !Ref EUCAccountIDs
 
   DataCollectionReadAccess:
     Type: AWS::IAM::ManagedPolicy