diff --git a/lambda_handler.py b/lambda_handler.py index 20d3489..8577c8b 100644 --- a/lambda_handler.py +++ b/lambda_handler.py @@ -1,6 +1,5 @@ -import os import json -from library import storage, aws, utility +from library import util, aws, storage from library.conf_gen import ConfigGenerator @@ -8,86 +7,30 @@ 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 = {} diff --git a/library/aws.py b/library/aws.py index 0b37302..5116042 100644 --- a/library/aws.py +++ b/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) \ No newline at end of file + 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 diff --git a/library/conf_gen.py b/library/conf_gen.py index c3be9ab..6692392 100644 --- a/library/conf_gen.py +++ b/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) \ No newline at end of file + 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 diff --git a/library/storage.py b/library/storage.py index a69e886..663a36b 100644 --- a/library/storage.py +++ b/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 \ No newline at end of file + 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") diff --git a/library/utility.py b/library/util.py similarity index 71% rename from library/utility.py rename to library/util.py index a3bcb43..1c6a4e3 100644 --- a/library/utility.py +++ b/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) \ No newline at end of file + 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' + }