Skip to content

Commit

Permalink
Use get_key in util file
Browse files Browse the repository at this point in the history
  • Loading branch information
@NoahStapp
NoahStapp committed Aug 29, 2023
1 parent 55819d1 commit d71795e
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 37 deletions.
48 changes: 21 additions & 27 deletions .evergreen/auth_aws/aws_tester.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import json
import sys
import subprocess
from util import get_key
from pymongo import MongoClient
from urllib.parse import quote_plus

Expand All @@ -17,7 +18,7 @@
from aws_assume_web_role import _assume_role_with_web_identity
from aws_assign_instance_profile import _assign_instance_policy

ASSUMED_ROLE = "arn:aws:sts::557821124784:assumed-role/authtest_user_assume_role/*";
ASSUMED_ROLE = "arn:aws:sts::557821124784:assumed-role/authtest_user_assume_role/*"
ASSUMED_WEB_ROLE = "arn:aws:sts::857654397073:assumed-role/webIdentityTestRole/*"

# This varies based on hosting EC2 as the account id and role name can vary
Expand All @@ -32,13 +33,6 @@
CONFIG = os.environ


def get_key(key: str) -> str:
if _USE_AWS_SECRETS:
return key.upper()
else:
return key


def run(args, env):
"""Run a python command in a subprocess."""
return subprocess.run([sys.executable] + args, env=env).returncode
Expand All @@ -60,10 +54,10 @@ def create_user(user, kwargs):

def setup_assume_role():
# Assume the role to get temp creds.
os.environ['AWS_ACCESS_KEY_ID'] = CONFIG[get_key("iam_auth_assume_aws_account")]
os.environ['AWS_SECRET_ACCESS_KEY'] = CONFIG[get_key("iam_auth_assume_aws_secret_access_key")]
os.environ['AWS_ACCESS_KEY_ID'] = CONFIG[get_key("iam_auth_assume_aws_account", _USE_AWS_SECRETS)]
os.environ['AWS_SECRET_ACCESS_KEY'] = CONFIG[get_key("iam_auth_assume_aws_secret_access_key", _USE_AWS_SECRETS)]

role_name = CONFIG[get_key("iam_auth_assume_role_name")]
role_name = CONFIG[get_key("iam_auth_assume_role_name", _USE_AWS_SECRETS)]
creds = _assume_role(role_name)
with open(os.path.join(HERE, 'creds.json'), 'w') as fid:
json.dump(creds, fid)
Expand All @@ -88,14 +82,14 @@ def setup_ecs():
mongo_binaries = os.environ['MONGODB_BINARIES']
project_dir = os.environ['PROJECT_DIRECTORY']
base_command = f"{sys.executable} -u lib/container_tester.py"
run_prune_command = f"{base_command} -v remote_gc_services --cluster {CONFIG[get_key('iam_auth_ecs_cluster')]}"
run_test_command = f"{base_command} -d -v run_e2e_test --cluster {CONFIG[get_key('iam_auth_ecs_cluster')]} --task_definition {CONFIG[get_key('iam_auth_ecs_task_definition')]} --subnets {CONFIG[get_key('iam_auth_ecs_subnet_a')]} --subnets {CONFIG[get_key('iam_auth_ecs_subnet_b')]} --security_group {CONFIG[get_key('iam_auth_ecs_security_group')]} --files {mongo_binaries}/mongod:/root/mongod {mongo_binaries}/mongosh:/root/mongosh lib/ecs_hosted_test.js:/root/ecs_hosted_test.js {project_dir}:/root --script lib/ecs_hosted_test.sh"
run_prune_command = f"{base_command} -v remote_gc_services --cluster {CONFIG[get_key('iam_auth_ecs_cluster', _USE_AWS_SECRETS)]}"
run_test_command = f"{base_command} -d -v run_e2e_test --cluster {CONFIG[get_key('iam_auth_ecs_cluster', _USE_AWS_SECRETS)]} --task_definition {CONFIG[get_key('iam_auth_ecs_task_definition', _USE_AWS_SECRETS)]} --subnets {CONFIG[get_key('iam_auth_ecs_subnet_a', _USE_AWS_SECRETS)]} --subnets {CONFIG[get_key('iam_auth_ecs_subnet_b', _USE_AWS_SECRETS)]} --security_group {CONFIG[get_key('iam_auth_ecs_security_group', _USE_AWS_SECRETS)]} --files {mongo_binaries}/mongod:/root/mongod {mongo_binaries}/mongosh:/root/mongosh lib/ecs_hosted_test.js:/root/ecs_hosted_test.js {project_dir}:/root --script lib/ecs_hosted_test.sh"

# Pass in the AWS credentials as environment variables
# AWS_SHARED_CREDENTIALS_FILE does not work in evergreen for an unknown
# reason
env = dict(AWS_ACCESS_KEY_ID=CONFIG[get_key('iam_auth_ecs_account')],
AWS_SECRET_ACCESS_KEY=CONFIG[get_key('iam_auth_ecs_secret_access_key')])
env = dict(AWS_ACCESS_KEY_ID=CONFIG[get_key('iam_auth_ecs_account', _USE_AWS_SECRETS)],
AWS_SECRET_ACCESS_KEY=CONFIG[get_key('iam_auth_ecs_secret_access_key', _USE_AWS_SECRETS)])

# Prune other containers
subprocess.check_call(['/bin/sh', '-c', run_prune_command], env=env)
Expand All @@ -107,39 +101,39 @@ def setup_ecs():
def setup_regular():
# Create the user.
kwargs = dict(
username=CONFIG[get_key("iam_auth_ecs_account")],
password=CONFIG[get_key("iam_auth_ecs_secret_access_key")]
username=CONFIG[get_key("iam_auth_ecs_account", _USE_AWS_SECRETS)],
password=CONFIG[get_key("iam_auth_ecs_secret_access_key", _USE_AWS_SECRETS)]
)
create_user(CONFIG[get_key("iam_auth_ecs_account_arn")], kwargs)
create_user(CONFIG[get_key("iam_auth_ecs_account_arn", _USE_AWS_SECRETS)], kwargs)


def setup_web_identity():
# Unassign the instance profile.
env = dict(AWS_ACCESS_KEY_ID=CONFIG[get_key("iam_auth_ec2_instance_account")],
AWS_SECRET_ACCESS_KEY=CONFIG[get_key("iam_auth_ec2_instance_secret_access_key")])
env = dict(AWS_ACCESS_KEY_ID=CONFIG[get_key("iam_auth_ec2_instance_account", _USE_AWS_SECRETS)],
AWS_SECRET_ACCESS_KEY=CONFIG[get_key("iam_auth_ec2_instance_secret_access_key", _USE_AWS_SECRETS)])
ret = run(['lib/aws_unassign_instance_profile.py'], env)
if ret == 2:
raise RuntimeError("Request limit exceeded for AWS API");
raise RuntimeError("Request limit exceeded for AWS API")

if ret != 0:
print('ret was', ret)
raise RuntimeError("Failed to unassign an instance profile from the current machine")

# Handle the OIDC credentials.
env = dict(
IDP_ISSUER=CONFIG[get_key("iam_web_identity_issuer")],
IDP_JWKS_URI=CONFIG[get_key("iam_web_identity_jwks_uri")],
IDP_RSA_KEY=CONFIG[get_key("iam_web_identity_rsa_key")],
AWS_WEB_IDENTITY_TOKEN_FILE=CONFIG[get_key('iam_web_identity_token_file')]
IDP_ISSUER=CONFIG[get_key("iam_web_identity_issuer", _USE_AWS_SECRETS)],
IDP_JWKS_URI=CONFIG[get_key("iam_web_identity_jwks_uri", _USE_AWS_SECRETS)],
IDP_RSA_KEY=CONFIG[get_key("iam_web_identity_rsa_key", _USE_AWS_SECRETS)],
AWS_WEB_IDENTITY_TOKEN_FILE=CONFIG[get_key('iam_web_identity_token_file', _USE_AWS_SECRETS)]
)

ret = run(['lib/aws_handle_oidc_creds.py', 'token'], env)
if ret != 0:
raise RuntimeWarning("Failed to write the web token")

# Assume the web role to get temp credentials.
os.environ['AWS_WEB_IDENTITY_TOKEN_FILE'] = CONFIG[get_key('iam_web_identity_token_file')]
os.environ['AWS_ROLE_ARN'] = CONFIG[get_key("iam_auth_assume_web_role_name")]
os.environ['AWS_WEB_IDENTITY_TOKEN_FILE'] = CONFIG[get_key('iam_web_identity_token_file', _USE_AWS_SECRETS)]
os.environ['AWS_ROLE_ARN'] = CONFIG[get_key("iam_auth_assume_web_role_name", _USE_AWS_SECRETS)]

creds = _assume_role_with_web_identity()
with open(os.path.join(HERE, 'creds.json'), 'w') as fid:
Expand Down
31 changes: 21 additions & 10 deletions .evergreen/auth_aws/lib/aws_assign_instance_profile.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,17 @@
import boto3
import botocore

sys.path.insert(1, os.path.join(sys.path[0], '..'))

from util import get_key

LOGGER = logging.getLogger(__name__)
HERE = os.path.abspath(os.path.dirname(__file__))


def _get_local_instance_id():
return urllib.request.urlopen('http://169.254.169.254/latest/meta-data/instance-id').read().decode()


def _has_instance_profile():
base_url = "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
try:
Expand All @@ -45,6 +49,7 @@ def _has_instance_profile():

return True


def _wait_instance_profile():
retry = 60
while not _has_instance_profile() and retry:
Expand All @@ -54,23 +59,30 @@ def _wait_instance_profile():
if retry == 0:
raise ValueError("Timeout on waiting for instance profile")


def _handle_config():
try:
with open(os.path.join(HERE, '..', 'aws_e2e_setup.json')) as fid:
CONFIG = json.load(fid)
use_secrets = False
except FileNotFoundError:
CONFIG = os.environ
use_secrets = True

os.environ.setdefault('AWS_ACCESS_KEY_ID', CONFIG['iam_auth_ec2_instance_account'])
os.environ.setdefault('AWS_SECRET_ACCESS_KEY', CONFIG['iam_auth_ec2_instance_secret_access_key'])
return CONFIG['iam_auth_ec2_instance_profile']
try:
os.environ.setdefault('AWS_ACCESS_KEY_ID', CONFIG[get_key('iam_auth_ec2_instance_account', use_secrets)])
os.environ.setdefault('AWS_SECRET_ACCESS_KEY',
CONFIG[get_key('iam_auth_ec2_instance_secret_access_key', use_secrets)])
return CONFIG[get_key('iam_auth_ec2_instance_profile', use_secrets)]
except Exception as e:
print(e)
return ''


DEFAULT_ARN = _handle_config()

def _assign_instance_policy(iam_instance_arn=DEFAULT_ARN):

def _assign_instance_policy(iam_instance_arn=DEFAULT_ARN):
if _has_instance_profile():
print("IMPORTANT: Found machine already has instance profile, skipping the assignment")
return
Expand All @@ -79,13 +91,13 @@ def _assign_instance_policy(iam_instance_arn=DEFAULT_ARN):

ec2_client = boto3.client("ec2", 'us-east-1')

#https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#EC2.Client.associate_iam_instance_profile
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#EC2.Client.associate_iam_instance_profile
try:
response = ec2_client.associate_iam_instance_profile(
IamInstanceProfile={
'Arn' : iam_instance_arn,
'Arn': iam_instance_arn,
},
InstanceId = instance_id)
InstanceId=instance_id)

print(response)

Expand All @@ -111,7 +123,6 @@ def main() -> None:

args = parser.parse_args()


if args.debug:
logging.basicConfig(level=logging.DEBUG)
elif args.verbose:
Expand All @@ -121,4 +132,4 @@ def main() -> None:


if __name__ == "__main__":
main()
main()
5 changes: 5 additions & 0 deletions .evergreen/auth_aws/util.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
def get_key(key: str, uppercase: bool) -> str:
if uppercase:
return key.upper()
else:
return key

0 comments on commit d71795e

Please sign in to comment.