Update a dotenv file with AWS credentials directly from aws-vault

aws-dotenv-update.py

#!/usr/bin/env python3
#
# AUTHOR: Jardel Weyrich <jweyrich at gmail dot com>
#
import argparse
import os
import subprocess

DEFAULT_DOTENV_FILE = '.env.dev'
DEFAULT_PROFILE = 'default'
DEFAULT_VAR_PREFIX = 'AWS_'

def get_awsenv_vars(profile, var_prefix):
    command = ['aws-vault', 'exec', profile, '--', 'env']
    # Retrieve AWS variables by executing `aws-vault` and reading its stdout
    result = subprocess.run(command, stdout=subprocess.PIPE)
    # Convert buffer b'k1=v1\nk2=v2\n' to string 'k1=v1\nk2=v2\n'
    data = result.stdout.decode('utf-8')
    # Convert 'k1=v1\nk2=v2\n' to ['k1=v1','k2=v2']
    all_envs = data.split(os.linesep)
    # Filter out non-AWS variables
    aws_envs = list(filter(lambda x: x.startswith(var_prefix), all_envs))
    # Convert ['k1=v1','k2=v2'] to {'k1':'v1','k2':'v2'}
    aws_vars = dict(v.split('=', 1) for v in aws_envs)
    return aws_vars

def read_dotenv_file(file):
    with open(file, "r") as f:
        data = f.read()
        # Convert 'k1=v1\nk2=v2\n' to ['k1=v1','k2=v2']
        lines = data.split(os.linesep)
        # Remove empty strings
        filtered_lines = filter(None, lines)
        # Convert ['k1=v1','k2=v2'] to {'k1':'v1','k2':'v2'}
        vars = dict(v.split('=', 1) for v in filtered_lines)
        return vars

def write_dotenv_file(file, vars):
    with open(file, "w") as f:
        # Convert {'k1':'v1','k2':'v2'} to ['k1=v1','k2=v2']
        vars_list = ['%s=%s' % (str(k), str(v)) for k, v in vars.items()]
        # Convert ['k1=v1','k2=v2'] to 'k1=v1\nk2=v2\n'
        data = os.linesep.join(vars_list) + os.linesep
        f.write(data)

def merge_dotenv_vars(dotenv_vars, vars_to_update):
    result = dict(dotenv_vars) # Clone it
    result.update(vars_to_update) # Update `result` dict in-place
    return result

def main():
    parser = argparse.ArgumentParser(description='Update only AWS credentials in the specified dotenv file.')
    parser.add_argument(
        '-p', '--profile', nargs='?',
        default=DEFAULT_PROFILE, help='AWS profile passed to aws-vault')
    parser.add_argument(
        '-v', '--var-prefix', nargs='?',
        default=DEFAULT_VAR_PREFIX, help='Filter variables using this prefix')
    parser.add_argument(
        '-d', '--dotenv-file', nargs='?',
        default=DEFAULT_DOTENV_FILE, help='Dotenv file to be updated')
    args = parser.parse_args()
    dotenv_vars = read_dotenv_file(args.dotenv_file)
    awsenv_vars = get_awsenv_vars(args.profile, args.var_prefix)
    merged_dotenv_vars = merge_dotenv_vars(dotenv_vars, awsenv_vars)
    write_dotenv_file(args.dotenv_file, merged_dotenv_vars)

if __name__ == "__main__":
    # execute only if run as a script
    main()