-
Notifications
You must be signed in to change notification settings - Fork 3
/
awsenv
executable file
·78 lines (59 loc) · 2.2 KB
/
awsenv
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
#!/usr/bin/env python3
import argparse
import os
import sys
from botocore.configloader import raw_config_parse
from botocore.session import Session
from botocore.exceptions import ClientError, InvalidConfigError
from botocore.exceptions import SSOTokenLoadError, UnauthorizedSSOTokenError
def get_profiles():
return Session().available_profiles
def get_credentials(profile):
AWS_CONF = '~/.aws/config'
AWS_VARS = ['AWS_ACCESS_KEY_ID',
'AWS_SECRET_ACCESS_KEY',
'AWS_SESSION_TOKEN']
try:
session = Session(profile=profile)
credentials = session.get_credentials().get_frozen_credentials()
except (SSOTokenLoadError, UnauthorizedSSOTokenError):
try:
os.system(f'aws --profile {profile} sso login')
return get_credentials(profile)
except Exception as e:
raise SystemExit(f'aws sso login failed with: {e.__class__.__name__}: {e}')
except Exception as e:
raise SystemExit(f'get_credentials failed with {e.__class__.__name__}: {e}')
else:
return list(zip(AWS_VARS, [(c or '') for c in credentials]))
def exec_sh(profile, cmd, creds):
NAME = os.path.basename(__file__)
GREEN = '\[$(tput setaf 2)\]'
RESET = '\[$(tput sgr0)\]'
os.environ['PS1'] = f'[\\u@{GREEN}{NAME}/{profile}{RESET} \\W]\\$ '
for k, v in creds:
os.environ[k] = str(v)
try:
os.execvp(cmd[0], cmd)
except OSError as e:
raise SystemExit(f"error executing '{cmd[0]}': {os.strerror(e.errno)}")
def process_args():
ap = argparse.ArgumentParser(
prog='awsenv',
description='run a command using a configured AWS profile',
usage='%(prog)s <profile> [--] [<command>]'
)
ap.add_argument('profile', choices=get_profiles(),
help='the AWS profile to use')
ap.add_argument('command', nargs=argparse.REMAINDER,
help='the command to execute')
return ap.parse_args()
def main():
args = process_args()
creds = get_credentials(args.profile)
if args.command:
exec_sh(args.profile, args.command, creds)
else:
for k, v in creds:
print(f'export {k}="{v}"')
main()