PYTHON-5157 Convert aws tests to use python scripts

.evergreen/auth_aws/aws_setup.sh

@@ -6,108 +6,33 @@
 #   . ./aws_setup.sh <test-name>
 #
 # Handles AWS credential setup and exports relevant environment variables.
-# Assumes you have already set up secrets.
+# Sets up secrets if they have not already been set up.
 set -eu
 
 SCRIPT_DIR=$(dirname ${BASH_SOURCE[0]})
 . $SCRIPT_DIR/../handle-paths.sh
 pushd $SCRIPT_DIR
 
-# Ensure that secrets have already been set up.
-if [ ! -f "secrets-export.sh" ]; then
-    echo "ERROR: please run './setup-secrets.sh' in this folder"
-fi
-
 # Activate the venv and source the secrets file.
 . ./activate-authawsvenv.sh
-source secrets-export.sh
 
-if [ "$1" == "web-identity" ]; then
-    export AWS_WEB_IDENTITY_TOKEN_FILE="./token_file.txt"
-fi
-
-# Handle the test setup if not using env variables.
-case $1 in
-    session-creds)
-        echo "Running aws_tester.py with assume-role"
-        # Set up credentials with assume-role to create user in MongoDB and write AWS credentials.
-        python aws_tester.py "assume-role"
-        ;;
-    env-creds)
-        echo "Running aws_tester.py with regular"
-        # Set up credentials with regular to create user in MongoDB and write AWS credentials.
-        python aws_tester.py "regular"
-        ;;
-    *)
-        python aws_tester.py "$1"
-        ;;
-esac
-
-# If this is ecs, exit now.
-if [ "$1" == "ecs" ]; then
-    exit 0
+# Ensure that secrets have already been set up.
+if [ ! -f "./secrets-export.sh" ]; then
+    bash ./setup-secrets.sh
 fi
 
-# Convenience functions.
-urlencode () {
-  python -c "import sys, urllib.parse as ulp; sys.stdout.write(ulp.quote_plus(sys.argv[1]))" "$1"
-}
+# Remove any AWS creds that might be set in the parent env.
+unset AWS_ACCESS_KEY_ID
+unset AWS_SECRET_ACCESS_KEY
+unset AWS_SESSION_TOKEN
 
-jsonkey () {
-    python -c  "import json,sys;sys.stdout.write(json.load(sys.stdin)[sys.argv[1]])" "$1" < ./creds.json
-}
+source ./secrets-export.sh
 
-# Handle extra vars based on auth type.
-USER=""
-case $1 in
-    assume-role)
-        USER=$(jsonkey AccessKeyId)
-        USER=$(urlencode "$USER")
-        PASS=$(jsonkey SecretAccessKey)
-        PASS=$(urlencode "$PASS")
-        SESSION_TOKEN=$(jsonkey SessionToken)
-        SESSION_TOKEN=$(urlencode "$SESSION_TOKEN")
-        ;;
-
-    session-creds)
-        AWS_ACCESS_KEY_ID=$(jsonkey AccessKeyId)
-        AWS_SECRET_ACCESS_KEY=$(jsonkey SecretAccessKey)
-        AWS_SESSION_TOKEN=$(jsonkey SessionToken)
-
-        export AWS_ACCESS_KEY_ID
-        export AWS_SECRET_ACCESS_KEY
-        export AWS_SESSION_TOKEN
-        ;;
-
-    web-identity)
-        export AWS_ROLE_ARN=$IAM_AUTH_ASSUME_WEB_ROLE_NAME
-        export AWS_WEB_IDENTITY_TOKEN_FILE="$SCRIPT_DIR/$AWS_WEB_IDENTITY_TOKEN_FILE"
-        ;;
-
-    regular)
-        USER=$(urlencode "${IAM_AUTH_ECS_ACCOUNT}")
-        PASS=$(urlencode "${IAM_AUTH_ECS_SECRET_ACCESS_KEY}")
-        ;;
-
-    env-creds)
-        export AWS_ACCESS_KEY_ID=$IAM_AUTH_ECS_ACCOUNT
-        export AWS_SECRET_ACCESS_KEY=$IAM_AUTH_ECS_SECRET_ACCESS_KEY
-        ;;
-esac
-
-# Handle the URI.
-if [ -n "$USER" ]; then
-    MONGODB_URI="mongodb://$USER:$PASS@localhost"
-    export USER
-    export PASS
-else
-    MONGODB_URI="mongodb://localhost"
-fi
-MONGODB_URI="${MONGODB_URI}/aws?authMechanism=MONGODB-AWS"
-if [[ -n ${SESSION_TOKEN:-} ]]; then
-    MONGODB_URI="${MONGODB_URI}&authMechanismProperties=AWS_SESSION_TOKEN:${SESSION_TOKEN}"
+if [ -f $SCRIPT_DIR/test-env.sh ]; then
+    rm $SCRIPT_DIR/test-env.sh
 fi
 
-export MONGODB_URI="$MONGODB_URI"
+python aws_tester.py "$1"
+source $SCRIPT_DIR/test-env.sh
 
 popd

.evergreen/auth_aws/aws_tester.py

@@ -5,23 +5,27 @@
 
 import argparse
 import json
+import logging
 import os
 import subprocess
 import sys
 from functools import partial
+from pathlib import Path
 from urllib.parse import quote_plus
 
 from pymongo import MongoClient
 from pymongo.errors import OperationFailure
 
-HERE = os.path.abspath(os.path.dirname(__file__))
+HERE = Path(__file__).absolute().parent
+LOGGER = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(message)s")
 
 
 def join(*parts):
     return os.path.join(*parts).replace(os.sep, "/")
 
 
-sys.path.insert(0, join(HERE, "lib"))
+sys.path.insert(0, str(HERE / "lib"))
 from aws_assign_instance_profile import _assign_instance_policy
 from aws_assume_role import _assume_role
 from aws_assume_web_role import _assume_role_with_web_identity
@@ -35,7 +39,7 @@ def join(*parts):
 _USE_AWS_SECRETS = False
 
 try:
-    with open(join(HERE, "aws_e2e_setup.json")) as fid:
+    with (HERE / "aws_e2e_setup.json").open() as fid:
         CONFIG = json.load(fid)
         get_key = partial(_get_key, uppercase=False)
 except FileNotFoundError:
@@ -51,7 +55,7 @@ def run(args, env):
 
 def create_user(user, kwargs):
     """Create a user and verify access."""
-    print("Creating user", user)
+    LOGGER.info("Creating user %s", user)
     client = MongoClient(username="bob", password="pwd123")
     db = client["$external"]
     try:
@@ -76,7 +80,7 @@ def setup_assume_role():
 
     role_name = CONFIG[get_key("iam_auth_assume_role_name")]
     creds = _assume_role(role_name, quiet=True)
-    with open(join(HERE, "creds.json"), "w") as fid:
+    with (HERE / "creds.json").open("w") as fid:
         json.dump(creds, fid)
 
     # Create the user.
@@ -87,6 +91,11 @@ def setup_assume_role():
         authmechanismproperties=f"AWS_SESSION_TOKEN:{token}",
     )
     create_user(ASSUMED_ROLE, kwargs)
+    return dict(
+        USER=kwargs["username"],
+        PASS=kwargs["password"],
+        SESSION_TOKEN=creds["SessionToken"],
+    )
 
 
 def setup_ec2():
@@ -95,6 +104,7 @@ def setup_ec2():
     os.environ.pop("AWS_ACCESS_KEY_ID", None)
     os.environ.pop("AWS_SECRET_ACCESS_KEY", None)
     create_user(AWS_ACCOUNT_ARN, dict())
+    return dict()
 
 
 def setup_ecs():
@@ -138,6 +148,8 @@ def setup_ecs():
     # Run the test in a container
     subprocess.check_call(["/bin/sh", "-c", run_test_command], env=env)
 
+    return dict()
+
 
 def setup_regular():
     # Create the user.
@@ -147,6 +159,8 @@ def setup_regular():
     )
     create_user(CONFIG[get_key("iam_auth_ecs_account_arn")], kwargs)
 
+    return dict(USER=kwargs["username"], PASS=kwargs["password"])
+
 
 def setup_web_identity():
     # Unassign the instance profile.
@@ -161,7 +175,7 @@ def setup_web_identity():
         raise RuntimeError("Request limit exceeded for AWS API")
 
     if ret != 0:
-        print("ret was", ret)
+        LOGGER.debug("return code was %s", ret)
         raise RuntimeError(
             "Failed to unassign an instance profile from the current machine"
         )
@@ -186,10 +200,11 @@ def setup_web_identity():
 
     # Assume the web role to get temp credentials.
     os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
-    os.environ["AWS_ROLE_ARN"] = CONFIG[get_key("iam_auth_assume_web_role_name")]
+    role_arn = CONFIG[get_key("iam_auth_assume_web_role_name")]
+    os.environ["AWS_ROLE_ARN"] = role_arn
 
     creds = _assume_role_with_web_identity(True)
-    with open(join(HERE, "creds.json"), "w") as fid:
+    with (HERE / "creds.json").open("w") as fid:
         json.dump(creds, fid)
 
     # Create the user.
@@ -201,12 +216,37 @@ def setup_web_identity():
     )
     create_user(ASSUMED_WEB_ROLE, kwargs)
 
+    return dict(AWS_WEB_IDENTITY_TOKEN_FILE=token_file, AWS_ROLE_ARN=role_arn)
+
+
+def handle_creds(creds: dict):
+    if "USER" in creds:
+        USER = quote_plus(creds.pop("USER"))
+        PASS = quote_plus(creds.pop("PASS"))
+        MONGODB_URI = f"mongodb://{USER}:{PASS}@localhost"
+    else:
+        MONGODB_URI = "mongodb://localhost"
+    MONGODB_URI = f"{MONGODB_URI}/aws?authMechanism=MONGODB-AWS"
+    if "SESSION_TOKEN" in creds:
+        SESSION_TOKEN = quote_plus(creds.pop("SESSION_TOKEN"))
+        MONGODB_URI = (
+            f"{MONGODB_URI}&authMechanismProperties=AWS_SESSION_TOKEN:{SESSION_TOKEN}"
+        )
+    with (HERE / "test-env.sh").open("w", newline="\n") as fid:
+        fid.write("#!/usr/bin/env bash\n\n")
+        fid.write("set +x\n")
+        for key, value in creds.items():
+            fid.write(f"export {key}={value}\n")
+        fid.write(f"export MONGODB_URI={MONGODB_URI}\n")
+
 
 def main():
     parser = argparse.ArgumentParser(description="MONGODB-AWS tester.")
     sub = parser.add_subparsers(title="Tester subcommands", help="sub-command help")
 
-    run_assume_role_cmd = sub.add_parser("assume-role", help="Assume role test")
+    run_assume_role_cmd = sub.add_parser(
+        "assume-role", aliases=["session-creds"], help="Assume role test"
+    )
     run_assume_role_cmd.set_defaults(func=setup_assume_role)
 
     run_ec2_cmd = sub.add_parser("ec2", help="EC2 test")
@@ -215,14 +255,20 @@ def main():
     run_ecs_cmd = sub.add_parser("ecs", help="ECS test")
     run_ecs_cmd.set_defaults(func=setup_ecs)
 
-    run_regular_cmd = sub.add_parser("regular", help="Regular credentials test")
+    run_regular_cmd = sub.add_parser(
+        "regular", aliases=["env-creds"], help="Regular credentials test"
+    )
     run_regular_cmd.set_defaults(func=setup_regular)
 
     run_web_identity_cmd = sub.add_parser("web-identity", help="Web identity test")
     run_web_identity_cmd.set_defaults(func=setup_web_identity)
 
     args = parser.parse_args()
-    args.func()
+    func_name = args.func.__name__.replace("setup_", "")
+    LOGGER.info("Running aws_tester.py with %s...", func_name)
+    creds = args.func()
+    handle_creds(creds)
+    LOGGER.info("Running aws_tester.py with %s... done.", func_name)
 
 
 if __name__ == "__main__":

.evergreen/auth_aws/lib/aws_assign_instance_profile.py

@@ -36,9 +36,9 @@ def _has_instance_profile():
         print("Reading: " + base_url)
         iam_role = urllib.request.urlopen(base_url).read().decode()
     except urllib.error.HTTPError as e:
-        print(e)
         if e.code == 404:
             return False
+        print(e)
         raise e
 
     try:

.evergreen/orchestration/drivers_orchestration.py

@@ -55,14 +55,13 @@ def get_options():
     parser.add_argument(
         "--version",
         default="latest",
-        help='The version to download (Required). Use "latest" to download '
+        help='The version to download. Use "latest" to download '
         "the newest available version (including release candidates).",
     )
     parser.add_argument(
         "--topology",
         choices=["standalone", "replica_set", "sharded_cluster"],
-        default="standalone",
-        help="The topology of the server deployment",
+        help="The topology of the server deployment (defaults to standalone unless another flag like load_balancer is set)",
     )
     parser.add_argument(
         "--auth", action="store_true", help="Whether to add authentication"
@@ -78,6 +77,9 @@ def get_options():
     other_group.add_argument(
         "--load-balancer", action="store_true", help="Whether to use a load balancer"
     )
+    other_group.add_argument(
+        "--auth-aws", action="store_true", help="Whether to use MONGODB-AWS auth"
+    )
     other_group.add_argument(
         "--skip-crypt-shared",
         action="store_true",
@@ -148,6 +150,11 @@ def get_options():
         opts.mongo_orchestration_home = DRIVERS_TOOLS / ".evergreen/orchestration"
     if opts.mongodb_binaries is None:
         opts.mongodb_binaries = DRIVERS_TOOLS / "mongodb/bin"
+    if not opts.topology and opts.load_balancer:
+        opts.topology = "sharded_cluster"
+    if opts.auth_aws:
+        opts.auth = True
+        opts.orchestration_file = "auth-aws.json"
     if opts.topology == "standalone" or not opts.topology:
         opts.topology = "server"
     if not opts.version:
@@ -395,6 +402,8 @@ def run(opts):
     MO_EXPANSION_YML.write_text(
         MO_EXPANSION_YML.read_text() + f'\nMONGODB_URI: "{uri}"'
     )
+    MO_EXPANSION_SH.touch()
+    MO_EXPANSION_SH.write_text(MO_EXPANSION_SH.read_text() + f'\nMONGODB_URI="{uri}"')
     URI_TXT.write_text(uri)
     LOGGER.info(f"Cluster URI: {uri}")
 

.evergreen/run-orchestration.sh

@@ -15,6 +15,7 @@ set -eu
 #   SKIP_CRYPT_SHARED      Set to a non-empty string to skip downloading crypt_shared
 #   MONGODB_BINARIES       Set the path to the MONGODB_BINARIES for mongo orchestration.
 #   LOAD_BALANCER          Set to a non-empty string to enable load balancer. Only supported for sharded clusters.
+#   AUTH_AWS               Set to a non-empty string to enable MONGODB-AWS authentication.
 #   PYTHON                 Set the Python binary to use.
 #   INSTALL_LEGACY_SHELL   Set to a non-empty string to install the legacy mongo shell.
 #   TLS_CERT_KEY_FILE      Set a .pem file to be used as the tlsCertificateKeyFile option in mongo-orchestration

.gitignore

@@ -88,6 +88,7 @@ evergreen_config_generator/dist
 *-expansion.yml
 secrets-export.sh
 token_file.txt
+test-env.sh
 
 # Virtual envs
 .venv