refact: reduce cognitive complexity

lambda_handler.py

@@ -1,93 +1,36 @@
-import os
 import json
-from library import storage, aws, utility
+from library import util, aws, storage
 from library.conf_gen import ConfigGenerator
 
 
 ou_colors = {}
 ou_overrides = {
  "ou-fmth-o9vy8tjo": "",
 }
-json_data = json.dumps(ou_overrides)
 
 def main_handler(event, context):
  config_gen = ConfigGenerator()
+ config = util.load_environment_config()
+ print(f"Running locally: {config['running_locally']}")
 
- s3_bucket = os.environ['S3_BUCKET']
- assume_role_name = os.environ['ASSUME_ROLE']
- role_names = os.environ['ROLE_NAMES'].split(',')
- session_name = os.environ['SESSION_NAME']
- payer_account_ids = os.environ['PAYER_ACCOUNT_IDS'].split(',')
- ou_overrides = json.loads(os.environ['OU_OVERRIDES']) 
-
- # Local testing
- running_locally = os.environ.get('RUNNING_LOCALLY', 'true').lower() == 'true'
- print(f"Running locally: {running_locally}")
-
- # Initialise dictionaries to hold aggregated configurations for each role
- awscli_configs = {}
- awscli_prefixed_configs = {}
- browser_plugin_configs = {}
-
- ou_overrides = utility.parse_ou_overrides()
-
- # Initialise account_ou_mapping before the loop
  account_ou_mapping = {}
+ configs = {role_name: {'awscli': '', 'awscli_prefixed': '', 'browser_plugin': ''} for role_name in config['role_names']}
 
- for payer_account_id in payer_account_ids:
- assume_role_arn = f"arn:aws:iam::{payer_account_id}:role/{assume_role_name}"
- credentials = aws.assume_role(assume_role_arn, session_name + payer_account_id)
- if not credentials:
- print(f"Failed to assume role {assume_role_arn}")
- continue
-
- accounts = aws.list_payer_accounts(credentials)
+ for payer_account_id in config['payer_account_ids']:
+ accounts = aws.process_accounts_for_payer(payer_account_id, config, account_ou_mapping, config_gen)
  if not accounts:
- print(f"Failed to list accounts for payer account {payer_account_id}.")
  continue
+ for role_name in config['role_names']:
+ role_configs = config_gen.generate_configs_for_role(role_name, accounts, account_ou_mapping, config_gen)
+ for key, value in role_configs.items():
+ configs[role_name][key] += value
 
- # Get OUs for this payer account and update the mapping
- new_ou_mapping = aws.get_organizational_units(credentials, ou_overrides)
- if new_ou_mapping:
- account_ou_mapping.update(new_ou_mapping)
-
- for role_name in role_names:
- # Generate configurations for each role
- print(f"Before calling generate_awscli_config: {type(accounts)}")
- new_awscli_config = config_gen.generate_awscli_config(accounts, role_name)
- new_awscli_prefixed_config = config_gen.generate_awscli_config(accounts, role_name, include_profile_prefix=True)
- new_browser_plugin_config = config_gen.generate_browser_plugin_config(accounts, account_ou_mapping, role_name)
-
- # Append new configurations to existing ones
- awscli_configs[role_name] = awscli_configs.get(role_name, "") + new_awscli_config
- awscli_prefixed_configs[role_name] = awscli_prefixed_configs.get(role_name, "") + new_awscli_prefixed_config
- browser_plugin_configs[role_name] = browser_plugin_configs.get(role_name, "") + new_browser_plugin_config
-
- # Decide where to write configurations after generating all of them
- if running_locally:
- # Write configurations to local directory for each role
- for role_name in role_names:
- if role_name in awscli_configs: # Check if configuration exists
- storage.write_to_local('.', role_name, awscli_configs[role_name], "awscli-config")
- if role_name in awscli_prefixed_configs:
- storage.write_to_local('.', role_name, awscli_prefixed_configs[role_name], "awscli-config-prefixed")
- if role_name in browser_plugin_configs:
- storage.write_to_local('.', role_name, browser_plugin_configs[role_name], "browser-plugin-config")
- else:
- # Write the generated configurations to S3
- for role_name in role_names:
- if role_name in awscli_configs: # Similar check as above
- storage.write_to_s3(s3_bucket, role_name, awscli_configs[role_name], "awscli-config")
- if role_name in awscli_prefixed_configs:
- storage.write_to_s3(s3_bucket, role_name, awscli_prefixed_configs[role_name], "awscli-config-prefixed")
- if role_name in browser_plugin_configs:
- storage.write_to_s3(s3_bucket, role_name, browser_plugin_configs[role_name], "browser-plugin-config")
-
+ storage.write_configs(config['role_names'], configs, config['running_locally'], config['s3_bucket'])
  print("Configuration files for all roles and payer accounts successfully processed.")
- # Export OU information after processing all roles and payer accounts
- utility.export_ou_information(account_ou_mapping, running_locally, s3_bucket, '.')
+ util.export_ou_information(account_ou_mapping, config['running_locally'], config['s3_bucket'], '.')
  print("OU information has been exported.")
 
+
 if __name__ == '__main__':
  mock_event = {"Update": "True"}
  mock_context = {}

library/aws.py

@@ -119,4 +119,22 @@ def process_ous(org_client, parent_id, account_ou_mapping, parent_name, ou_overr
  account_ou_mapping[account_id]['OUId'] = ou['Id']
 
  # Recursively process any child OUs
- process_ous(org_client, ou['Id'], account_ou_mapping, new_parent_name, ou_overrides)
+ process_ous(org_client, ou['Id'], account_ou_mapping, new_parent_name, ou_overrides)
+
+def process_accounts_for_payer(payer_account_id, config, account_ou_mapping, config_gen):
+ assume_role_arn = f"arn:aws:iam::{payer_account_id}:role/{config['assume_role_name']}"
+ credentials = assume_role(assume_role_arn, config['session_name'] + payer_account_id)
+ if not credentials:
+ print(f"Failed to assume role {assume_role_arn}")
+ return
+
+ accounts = list_payer_accounts(credentials)
+ if not accounts:
+ print(f"Failed to list accounts for payer account {payer_account_id}.")
+ return
+
+ new_ou_mapping = get_organizational_units(credentials, config['ou_overrides'])
+ if new_ou_mapping:
+ account_ou_mapping.update(new_ou_mapping)
+
+ return accounts

library/conf_gen.py

@@ -72,4 +72,12 @@ def generate_browser_plugin_config(self, accounts, account_ou_mapping, role_name
  config_lines.append(f"color = {color}") # Include the color in the config
  config_lines.append("") # Add a newline for readability
 
- return "\n".join(config_lines)
+ return "\n".join(config_lines)
+
+ def generate_configs_for_role(self, role_name, accounts, account_ou_mapping, config_gen):
+ configs = {
+ 'awscli': config_gen.generate_awscli_config(accounts, role_name),
+ 'awscli_prefixed': config_gen.generate_awscli_config(accounts, role_name, include_profile_prefix=True),
+ 'browser_plugin': config_gen.generate_browser_plugin_config(accounts, account_ou_mapping, role_name)
+ }
+ return configs

library/storage.py

@@ -115,4 +115,15 @@ def write_ou_info_to_s3(bucket_name, ou_info_content):
  return True
  except Exception as e:
  print(f"Error uploading OU information to S3: {e}")
- return False
+ return False
+
+def write_configs(role_names, configs, running_locally, s3_bucket):
+ for role_name in role_names:
+ if running_locally:
+ write_to_local('.', role_name, configs[role_name]['awscli'], "awscli-config")
+ write_to_local('.', role_name, configs[role_name]['awscli_prefixed'], "awscli-config-prefixed")
+ write_to_local('.', role_name, configs[role_name]['browser_plugin'], "browser-plugin-config")
+ else:
+ write_to_s3(s3_bucket, role_name, configs[role_name]['awscli'], "awscli-config")
+ write_to_s3(s3_bucket, role_name, configs[role_name]['awscli_prefixed'], "awscli-config-prefixed")
+ write_to_s3(s3_bucket, role_name, configs[role_name]['browser_plugin'], "browser-plugin-config")

library/utility.py → library/util.py

@@ -33,4 +33,15 @@ def export_ou_information(ou_mapping, running_locally, s3_bucket=None, directory
  if running_locally:
  storage.write_ou_info_to_local(directory, ou_info_content)
  else:
- storage.write_ou_info_to_s3(s3_bucket, ou_info_content)
+ storage.write_ou_info_to_s3(s3_bucket, ou_info_content)
+
+def load_environment_config():
+ return {
+ 's3_bucket': os.environ['S3_BUCKET'],
+ 'assume_role_name': os.environ['ASSUME_ROLE'],
+ 'role_names': os.environ['ROLE_NAMES'].split(','),
+ 'session_name': os.environ['SESSION_NAME'],
+ 'payer_account_ids': os.environ['PAYER_ACCOUNT_IDS'].split(','),
+ 'ou_overrides': json.loads(os.environ['OU_OVERRIDES']),
+ 'running_locally': os.environ.get('RUNNING_LOCALLY', 'true').lower() == 'true'
+ }