Skip to content

Commit ba75f25

Browse files
blink1073eramongodb
blink1073
and
eramongodb
authored
PYTHON-5157 Convert aws tests to use python scripts (#624)
Co-authored-by: Ezra Chung <88335979+eramongodb@users.noreply.github.com>
1 parent cf90abd commit ba75f25

8 files changed

+194
-110
lines changed

‎.evergreen/auth_aws/aws_setup.sh

+13-88
Original file line numberDiff line numberDiff line change
@@ -6,108 +6,33 @@
66
# . ./aws_setup.sh <test-name>
77
#
88
# Handles AWS credential setup and exports relevant environment variables.
9-
# Assumes you have already set up secrets.
9+
# Sets up secrets if they have not already been set up.
1010
set -eu
1111

1212
SCRIPT_DIR=$(dirname ${BASH_SOURCE[0]})
1313
. $SCRIPT_DIR/../handle-paths.sh
1414
pushd $SCRIPT_DIR
1515

16-
# Ensure that secrets have already been set up.
17-
if [ ! -f "secrets-export.sh" ]; then
18-
echo "ERROR: please run './setup-secrets.sh' in this folder"
19-
fi
20-
2116
# Activate the venv and source the secrets file.
2217
. ./activate-authawsvenv.sh
23-
source secrets-export.sh
2418

25-
if [ "$1" == "web-identity" ]; then
26-
export AWS_WEB_IDENTITY_TOKEN_FILE="./token_file.txt"
27-
fi
28-
29-
# Handle the test setup if not using env variables.
30-
case $1 in
31-
session-creds)
32-
echo "Running aws_tester.py with assume-role"
33-
# Set up credentials with assume-role to create user in MongoDB and write AWS credentials.
34-
python aws_tester.py "assume-role"
35-
;;
36-
env-creds)
37-
echo "Running aws_tester.py with regular"
38-
# Set up credentials with regular to create user in MongoDB and write AWS credentials.
39-
python aws_tester.py "regular"
40-
;;
41-
*)
42-
python aws_tester.py "$1"
43-
;;
44-
esac
45-
46-
# If this is ecs, exit now.
47-
if [ "$1" == "ecs" ]; then
48-
exit 0
19+
# Ensure that secrets have already been set up.
20+
if [ ! -f "./secrets-export.sh" ]; then
21+
bash ./setup-secrets.sh
4922
fi
5023

51-
# Convenience functions.
52-
urlencode () {
53-
python -c "import sys, urllib.parse as ulp; sys.stdout.write(ulp.quote_plus(sys.argv[1]))" "$1"
54-
}
24+
# Remove any AWS creds that might be set in the parent env.
25+
unset AWS_ACCESS_KEY_ID
26+
unset AWS_SECRET_ACCESS_KEY
27+
unset AWS_SESSION_TOKEN
5528

56-
jsonkey () {
57-
python -c "import json,sys;sys.stdout.write(json.load(sys.stdin)[sys.argv[1]])" "$1" < ./creds.json
58-
}
29+
source ./secrets-export.sh
5930

60-
# Handle extra vars based on auth type.
61-
USER=""
62-
case $1 in
63-
assume-role)
64-
USER=$(jsonkey AccessKeyId)
65-
USER=$(urlencode "$USER")
66-
PASS=$(jsonkey SecretAccessKey)
67-
PASS=$(urlencode "$PASS")
68-
SESSION_TOKEN=$(jsonkey SessionToken)
69-
SESSION_TOKEN=$(urlencode "$SESSION_TOKEN")
70-
;;
71-
72-
session-creds)
73-
AWS_ACCESS_KEY_ID=$(jsonkey AccessKeyId)
74-
AWS_SECRET_ACCESS_KEY=$(jsonkey SecretAccessKey)
75-
AWS_SESSION_TOKEN=$(jsonkey SessionToken)
76-
77-
export AWS_ACCESS_KEY_ID
78-
export AWS_SECRET_ACCESS_KEY
79-
export AWS_SESSION_TOKEN
80-
;;
81-
82-
web-identity)
83-
export AWS_ROLE_ARN=$IAM_AUTH_ASSUME_WEB_ROLE_NAME
84-
export AWS_WEB_IDENTITY_TOKEN_FILE="$SCRIPT_DIR/$AWS_WEB_IDENTITY_TOKEN_FILE"
85-
;;
86-
87-
regular)
88-
USER=$(urlencode "${IAM_AUTH_ECS_ACCOUNT}")
89-
PASS=$(urlencode "${IAM_AUTH_ECS_SECRET_ACCESS_KEY}")
90-
;;
91-
92-
env-creds)
93-
export AWS_ACCESS_KEY_ID=$IAM_AUTH_ECS_ACCOUNT
94-
export AWS_SECRET_ACCESS_KEY=$IAM_AUTH_ECS_SECRET_ACCESS_KEY
95-
;;
96-
esac
97-
98-
# Handle the URI.
99-
if [ -n "$USER" ]; then
100-
MONGODB_URI="mongodb://$USER:$PASS@localhost"
101-
export USER
102-
export PASS
103-
else
104-
MONGODB_URI="mongodb://localhost"
105-
fi
106-
MONGODB_URI="${MONGODB_URI}/aws?authMechanism=MONGODB-AWS"
107-
if [[ -n ${SESSION_TOKEN:-} ]]; then
108-
MONGODB_URI="${MONGODB_URI}&authMechanismProperties=AWS_SESSION_TOKEN:${SESSION_TOKEN}"
31+
if [ -f $SCRIPT_DIR/test-env.sh ]; then
32+
rm $SCRIPT_DIR/test-env.sh
10933
fi
11034

111-
export MONGODB_URI="$MONGODB_URI"
35+
python aws_tester.py "$@"
36+
source $SCRIPT_DIR/test-env.sh
11237

11338
popd

‎.evergreen/auth_aws/aws_tester.py

+78-9
Original file line numberDiff line numberDiff line change
@@ -5,23 +5,27 @@
55

66
import argparse
77
import json
8+
import logging
89
import os
910
import subprocess
1011
import sys
1112
from functools import partial
13+
from pathlib import Path
1214
from urllib.parse import quote_plus
1315

1416
from pymongo import MongoClient
1517
from pymongo.errors import OperationFailure
1618

17-
HERE = os.path.abspath(os.path.dirname(__file__))
19+
HERE = Path(__file__).absolute().parent
20+
LOGGER = logging.getLogger(__name__)
21+
logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(message)s")
1822

1923

2024
def join(*parts):
2125
return os.path.join(*parts).replace(os.sep, "/")
2226

2327

24-
sys.path.insert(0, join(HERE, "lib"))
28+
sys.path.insert(0, str(HERE / "lib"))
2529
from aws_assign_instance_profile import _assign_instance_policy
2630
from aws_assume_role import _assume_role
2731
from aws_assume_web_role import _assume_role_with_web_identity
@@ -35,7 +39,7 @@ def join(*parts):
3539
_USE_AWS_SECRETS = False
3640

3741
try:
38-
with open(join(HERE, "aws_e2e_setup.json")) as fid:
42+
with (HERE / "aws_e2e_setup.json").open() as fid:
3943
CONFIG = json.load(fid)
4044
get_key = partial(_get_key, uppercase=False)
4145
except FileNotFoundError:
@@ -51,7 +55,7 @@ def run(args, env):
5155

5256
def create_user(user, kwargs):
5357
"""Create a user and verify access."""
54-
print("Creating user", user)
58+
LOGGER.info("Creating user %s", user)
5559
client = MongoClient(username="bob", password="pwd123")
5660
db = client["$external"]
5761
try:
@@ -76,7 +80,7 @@ def setup_assume_role():
7680

7781
role_name = CONFIG[get_key("iam_auth_assume_role_name")]
7882
creds = _assume_role(role_name, quiet=True)
79-
with open(join(HERE, "creds.json"), "w") as fid:
83+
with (HERE / "creds.json").open("w") as fid:
8084
json.dump(creds, fid)
8185

8286
# Create the user.
@@ -87,6 +91,11 @@ def setup_assume_role():
8791
authmechanismproperties=f"AWS_SESSION_TOKEN:{token}",
8892
)
8993
create_user(ASSUMED_ROLE, kwargs)
94+
return dict(
95+
USER=kwargs["username"],
96+
PASS=kwargs["password"],
97+
SESSION_TOKEN=creds["SessionToken"],
98+
)
9099

91100

92101
def setup_ec2():
@@ -95,6 +104,7 @@ def setup_ec2():
95104
os.environ.pop("AWS_ACCESS_KEY_ID", None)
96105
os.environ.pop("AWS_SECRET_ACCESS_KEY", None)
97106
create_user(AWS_ACCOUNT_ARN, dict())
107+
return dict()
98108

99109

100110
def setup_ecs():
@@ -138,6 +148,18 @@ def setup_ecs():
138148
# Run the test in a container
139149
subprocess.check_call(["/bin/sh", "-c", run_test_command], env=env)
140150

151+
return dict()
152+
153+
154+
def setup_session_creds():
155+
# Set up the assume role user, and export the aws vars.
156+
creds = setup_assume_role()
157+
return dict(
158+
AWS_ACCESS_KEY_ID=creds["USER"],
159+
AWS_SECRET_ACCESS_KEY=creds["PASS"],
160+
AWS_SESSION_TOKEN=creds["SESSION_TOKEN"],
161+
)
162+
141163

142164
def setup_regular():
143165
# Create the user.
@@ -147,6 +169,14 @@ def setup_regular():
147169
)
148170
create_user(CONFIG[get_key("iam_auth_ecs_account_arn")], kwargs)
149171

172+
return dict(USER=kwargs["username"], PASS=kwargs["password"])
173+
174+
175+
def setup_env_creds():
176+
# Set up the regular user, but export the creds as environment vars.
177+
creds = setup_regular()
178+
return dict(AWS_ACCESS_KEY_ID=creds["USER"], AWS_SECRET_ACCESS_KEY=creds["PASS"])
179+
150180

151181
def setup_web_identity():
152182
# Unassign the instance profile.
@@ -161,7 +191,7 @@ def setup_web_identity():
161191
raise RuntimeError("Request limit exceeded for AWS API")
162192

163193
if ret != 0:
164-
print("ret was", ret)
194+
LOGGER.debug("return code was %s", ret)
165195
raise RuntimeError(
166196
"Failed to unassign an instance profile from the current machine"
167197
)
@@ -186,10 +216,11 @@ def setup_web_identity():
186216

187217
# Assume the web role to get temp credentials.
188218
os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
189-
os.environ["AWS_ROLE_ARN"] = CONFIG[get_key("iam_auth_assume_web_role_name")]
219+
role_arn = CONFIG[get_key("iam_auth_assume_web_role_name")]
220+
os.environ["AWS_ROLE_ARN"] = role_arn
190221

191222
creds = _assume_role_with_web_identity(True)
192-
with open(join(HERE, "creds.json"), "w") as fid:
223+
with (HERE / "creds.json").open("w") as fid:
193224
json.dump(creds, fid)
194225

195226
# Create the user.
@@ -201,6 +232,34 @@ def setup_web_identity():
201232
)
202233
create_user(ASSUMED_WEB_ROLE, kwargs)
203234

235+
return dict(AWS_WEB_IDENTITY_TOKEN_FILE=token_file, AWS_ROLE_ARN=role_arn)
236+
237+
238+
def handle_creds(creds: dict):
239+
if "USER" in creds:
240+
USER = quote_plus(creds["USER"])
241+
if "PASS" in creds:
242+
PASS = quote_plus(creds["PASS"])
243+
MONGODB_URI = f"mongodb://{USER}:{PASS}@localhost"
244+
else:
245+
MONGODB_URI = f"mongodb://{USER}@localhost"
246+
else:
247+
MONGODB_URI = "mongodb://localhost"
248+
MONGODB_URI = f"{MONGODB_URI}/aws?authMechanism=MONGODB-AWS"
249+
if "SESSION_TOKEN" in creds:
250+
SESSION_TOKEN = quote_plus(creds["SESSION_TOKEN"])
251+
MONGODB_URI = (
252+
f"{MONGODB_URI}&authMechanismProperties=AWS_SESSION_TOKEN:{SESSION_TOKEN}"
253+
)
254+
with (HERE / "test-env.sh").open("w", newline="\n") as fid:
255+
fid.write("#!/usr/bin/env bash\n\n")
256+
fid.write("set +x\n")
257+
for key, value in creds.items():
258+
if key in ["USER", "PASS", "SESSION_TOKEN"]:
259+
value = quote_plus(value) # noqa: PLW2901
260+
fid.write(f"export {key}={value}\n")
261+
fid.write(f"export MONGODB_URI={MONGODB_URI}\n")
262+
204263

205264
def main():
206265
parser = argparse.ArgumentParser(description="MONGODB-AWS tester.")
@@ -218,11 +277,21 @@ def main():
218277
run_regular_cmd = sub.add_parser("regular", help="Regular credentials test")
219278
run_regular_cmd.set_defaults(func=setup_regular)
220279

280+
run_session_creds_cmd = sub.add_parser("session-creds", help="Session credentials")
281+
run_session_creds_cmd.set_defaults(func=setup_session_creds)
282+
283+
run_env_creds_cmd = sub.add_parser("env-creds", help="Environment credentials")
284+
run_env_creds_cmd.set_defaults(func=setup_env_creds)
285+
221286
run_web_identity_cmd = sub.add_parser("web-identity", help="Web identity test")
222287
run_web_identity_cmd.set_defaults(func=setup_web_identity)
223288

224289
args = parser.parse_args()
225-
args.func()
290+
func_name = args.func.__name__.replace("setup_", "").replace("_", "-")
291+
LOGGER.info("Running aws_tester.py with %s...", func_name)
292+
creds = args.func()
293+
handle_creds(creds)
294+
LOGGER.info("Running aws_tester.py with %s... done.", func_name)
226295

227296

228297
if __name__ == "__main__":

‎.evergreen/auth_aws/lib/aws_assign_instance_profile.py

+9-9
Original file line numberDiff line numberDiff line change
@@ -33,23 +33,23 @@ def _get_local_instance_id():
3333
def _has_instance_profile():
3434
base_url = "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
3535
try:
36-
print("Reading: " + base_url)
36+
LOGGER.info("Reading: " + base_url)
3737
iam_role = urllib.request.urlopen(base_url).read().decode()
3838
except urllib.error.HTTPError as e:
39-
print(e)
4039
if e.code == 404:
4140
return False
41+
LOGGER.error(e)
4242
raise e
4343

4444
try:
4545
url = base_url + iam_role
46-
print("Reading: " + url)
46+
LOGGER.info("Reading: " + url)
4747
_ = urllib.request.urlopen(url)
48-
print("Assigned " + iam_role)
48+
LOGGER.info("Assigned " + iam_role)
4949
except urllib.error.HTTPError as e:
50-
print(e)
5150
if e.code == 404:
5251
return False
52+
LOGGER.error(e)
5353
raise e
5454

5555
return True
@@ -85,7 +85,7 @@ def _handle_config():
8585
)
8686
return CONFIG[get_key("iam_auth_ec2_instance_profile")]
8787
except Exception as e:
88-
print(e)
88+
LOGGER.error(e)
8989
return ""
9090

9191

@@ -94,7 +94,7 @@ def _handle_config():
9494

9595
def _assign_instance_policy(iam_instance_arn=DEFAULT_ARN):
9696
if _has_instance_profile():
97-
print(
97+
LOGGER.warning(
9898
"IMPORTANT: Found machine already has instance profile, skipping the assignment"
9999
)
100100
return
@@ -112,14 +112,14 @@ def _assign_instance_policy(iam_instance_arn=DEFAULT_ARN):
112112
InstanceId=instance_id,
113113
)
114114

115-
print(response)
115+
LOGGER.debug(response)
116116

117117
# Wait for the instance profile to be assigned by polling the local instance metadata service
118118
_wait_instance_profile()
119119

120120
except botocore.exceptions.ClientError as ce:
121121
if ce.response["Error"]["Code"] == "RequestLimitExceeded":
122-
print("WARNING: RequestLimitExceeded, exiting with error code 2")
122+
LOGGER.warning("WARNING: RequestLimitExceeded, exiting with error code 2")
123123
sys.exit(2)
124124
raise
125125

0 commit comments

Comments
 (0)