#!/usr/bin/env python3

import click
import subprocess
import configparser
import os
import json
import io

TB_SCRIPT = """
tell application "Tunnelblick"
	set connected_configs to name of every configuration whose state = "CONNECTED"
	return connected_configs
end tell
"""

@click.command()
@click.option("--profile", default='default', type=str, help='AWS profile to update')
@click.option("--og", is_flag=True, help="Use OG method.")
@click.option("--from-profile/--not-from-profile", default=True, is_flag=True, help="Gather arguments from the aws profile.")
@click.option("--vault-path", default='aws/creds', type=str, help="Default path in the vault for keys.")
@click.option("--vault-write/--vault-read", default=True, is_flag=True, help="Should we read or write from the vault?")
@click.option("--role", "--vault-role", default='user_engineer_default', type=str)
@click.option("--config", "config_path", default=os.path.expanduser('~/.aws/credentials'), type=click.Path())
def main(profile, og, from_profile, vault_path, vault_write, vault_role, config_path):
    """Update AWS credentials from the LendUp Vault"""
    if not check_vpn():
        click.echo("Please connect to the prod-us-east VPN before continuing")
        raise click.Abort()
    config = get_aws_config(config_path)
    
    if from_profile and profile in config:
        vault_role = config[profile].get('vault_role', vault_role)
        vault_path = config[profile].get('vault_path', vault_path)
        vault_write = config[profile].get('vault_write', vault_write)
        
    if og:
        vault_path = "aws/sts"
        vault_write = True
        if not vault_role.endswith('-og'):
            vault_role += '-og'
    
    try:
        response_data = get_vault_info(path=vault_path, role=vault_role, write=vault_write)
    except subprocess.CalledProcessError as e:
        click.echo("Error in vault read, returned exit code: {}".format(e.returncode))
        raise click.Abort()
    if profile not in config:
        config.add_section(profile)
    
    if from_profile:
        config[profile]['vault_role'] = vault_role
        config[profile]['vault_path'] = vault_path
        config[profile]['vault_write'] = vault_write
    
    config[profile]['aws_access_key_id'] = response_data['access_key']
    config[profile]['aws_secret_access_key'] = response_data['secret_key']
    if 'security_token' in response_data:
        config[profile]['aws_session_token'] = response_data['security_token']
    else:
        config[profile].pop('aws_session_token')
    
    with open(config_path, 'w') as f:
        config.write(f)

def get_vault_info(path, role, write):
    """Get and parse fixed-width format Vault information."""
    if write:
        out = subprocess.check_output(['vault', 'write', '-f', '--format', 'json', f'{path}{role}'])
    else:
        out = subprocess.check_output(['vault', 'read', '--format', 'json', f'{path}{role}'])
    response = json.loads(out)
    return response['data']

def check_vpn(connection='prod-us-east'):
    """Check VPN Connection"""
    out = subprocess.check_output(["osascript", "-e", TB_SCRIPT]).decode('utf-8')
    return connection in out
    
def get_aws_config(path="~/.aws/credentials"):
    """Get the AWS configuration"""
    config = configparser.ConfigParser()
    config.read(path)
    return config
    
if __name__ == '__main__':
    main()
@@ -78,16 +78,33 @@
 port os%0A
+import posixpath%0A
 import j
@@ -302,16 +302,117 @@
 mmand()%0A
+@click.option(%22--auth/--no-auth%22, default=True, is_flag=True, help=%22Authenticate with vault first.%22)%0A
 @click.o
@@ -1092,16 +1092,78 @@
 Path())%0A
+@click.option(%22--vault-username%22, default=os.environ%5B'USER'%5D)%0A
 def main
@@ -1163,16 +1163,22 @@
 ef main(
+auth, 
 profile,
@@ -1244,16 +1244,32 @@
 fig_path
+, vault_username
 ):%0A    %22
@@ -1445,24 +1445,291 @@
 ick.Abort()%0A
+    if not check_vault_connection():%0A        if auth:%0A            authenticate_vault(vault_username)%0A        else:%0A            click.echo(%22Please authenticate in vault with %60vault auth%60 or the %60--auth%60 flag to this command.%22)%0A            raise click.Abort()%0A        %0A
     config =
@@ -2660,24 +2660,33 @@
 lt_write'%5D =
+ %22yes%22 if
  vault_write
@@ -2685,16 +2685,26 @@
 lt_write
+ else %22no%22
 %0A    %0A  
@@ -3192,24 +3192,66 @@
 rmation.%22%22%22%0A
+    fullpath = posixpath.join(path, role)%0A
     if write
@@ -3334,34 +3334,32 @@
  'json', f'%7B
+full
 path
-%7D%7Brole
 %7D'%5D)%0A    els
@@ -3444,18 +3444,16 @@
  f'%7B
+full
 path
-%7D%7Brole
 %7D'%5D)
@@ -3511,16 +3511,516 @@
 %5B'data'%5D
+%0A    %0Adef check_vault_connection():%0A    %22%22%22docstring for check_vault_connection%22%22%22%0A    try:%0A        out = subprocess.check_output(%5B'vault', 'read', '--format', 'json', '/auth/token/lookup-self'%5D) == 0%0A    except subprocess.CalledProcessError as e:%0A        return False%0A    else:%0A        return True%0A%0Adef authenticate_vault(username):%0A    %22%22%22Authenticate to Vault.%22%22%22%0A    args = %5B'vault', 'auth', '-method=ldap', 'username=%7B%7D'.format(username)%5D%0A    click.echo(%22 %22.join(args))%0A    subprocess.call(args)
 %0A%0Adef ch

Update vault script