-
-
Notifications
You must be signed in to change notification settings - Fork 13
/
aws_env.py
executable file
·93 lines (80 loc) · 3.02 KB
/
aws_env.py
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
#! /usr/bin/env python3
"""Set AWS Environment variables from aws cli profiles."""
from __future__ import annotations
import os
import re
from typing import Dict, List, Optional
from utils import choose_from_list, run_cmd
def aws_version() -> Optional[int]:
"""Test if the aws cli version 2 is installed."""
try:
result = run_cmd(["aws", "--version"], check_results=False, chomp=True)
except FileNotFoundError:
print("AWS CLI version 2 is not installed.")
return None
else:
if result.returncode == 0:
# get major version number from stdout
match = re.match(r"aws-cli/(\d+)\..*", result.stdout)
if match:
return int(match.group(1))
return None
def list_aws_profiles() -> List[str]:
aws_ver = aws_version()
if aws_ver is not None and aws_ver == 2:
result = run_cmd(["aws", "configure", "list-profiles"], chomp=True)
return result.stdout.split("\n")
return []
def get_profile_var(profile: str, var_name: str) -> str:
result = run_cmd(["aws", "configure", "--profile", profile, "get", var_name], chomp=True)
return result.stdout
def init_aws_environment() -> None:
profile_list = list_aws_profiles()
# Build a map for looking up a profile name from the access key id. This
# algorithm assumes:
# - the 'default' profile will be processed first
# - if there are profiles using the same access key id, only the last
# one will be put into the map
if len(profile_list) == 0:
return
profile_map: Dict[str, str] = {}
for profile in profile_list:
key_id = get_profile_var(profile, "aws_access_key_id")
profile_map[key_id] = profile
curr_access_key = os.getenv("AWS_ACCESS_KEY_ID", "")
if curr_access_key in profile_map:
curr_profile = profile_map[curr_access_key]
else:
curr_profile = None
aws_profile = choose_from_list("AWS Environment", curr_profile, profile_list)
if aws_profile is not None and aws_profile != curr_profile:
os.environ["AWS_PROFILE"] = aws_profile
os.environ["AWS_ACCESS_KEY_ID"] = get_profile_var(aws_profile, "aws_access_key_id")
os.environ["AWS_SECRET_ACCESS_KEY"] = get_profile_var(aws_profile, "aws_secret_access_key")
region = get_profile_var(aws_profile, "region")
os.environ["AWS_DEFAULT_REGION"] = region
result = run_cmd(
[
"aws",
"sts",
"--profile",
aws_profile,
"get-caller-identity",
"--query",
"Account",
"--output",
"text",
],
chomp=True,
)
os.environ["AWS_ACCOUNT"] = result.stdout
if __name__ == "__main__":
init_aws_environment()
print("AWS Environment:")
for env_var in [
"AWS_ACCOUNT",
"AWS_DEFAULT_REGION",
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
]:
print(f"{env_var}: {os.getenv(env_var, None)}")