DRIVERS-2328 Add aws_setup convenience script (#365)

* GODRIVER-2607 Add aws_setup script

* cleanup

* fix handling of directory

* get the right secrets

* fix unbound var

* clean up script

* more cleanup

* fix handling of file

* activate venv

* allow user to already exist

* fix import

* fixups

* try again

* ignore windows for now

* debug

* debug

* move secrets handling out

* fix typo

* debug

* use robust script dir

* fix handling of token file

* syntax

* try adding os env

* try this way

.evergreen/auth_aws/aws_setup.sh

@@ -0,0 +1,105 @@
+#!/usr/bin/env bash
+#
+# aws_setup.sh
+#
+# Usage:
+# . ./aws_setup.sh <test-name>
+#
+# Handles AWS credential setup and exports relevant environment variables.
+# Assumes you have already set up secrets.
+set -eux
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+pushd $SCRIPT_DIR
+
+# Ensure that secrets have already been set up.
+if [ ! -f "secrets-export.sh" ]; then 
+ echo "ERROR: please run './setup_secrets.sh drivers/aws_auth' 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 | env-creds)
+ echo "Skipping aws_tester.py"
+ ;;
+ *)
+ python aws_tester.py "$1"
+ ;;
+esac
+
+# If this is ecs, exit now.
+if [ "$1" == "ecs" ]; then
+ exit 0
+fi
+
+# Convenience functions.
+urlencode () {
+ python -c "import sys, urllib.parse as ulp; sys.stdout.write(ulp.quote_plus(sys.argv[1]))" "$1"
+}
+
+jsonkey () {
+ python -c "import json,sys;sys.stdout.write(json.load(sys.stdin)[sys.argv[1]])" "$1" < ./creds.json
+}
+
+# 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}"
+fi
+
+export MONGODB_URI="$MONGODB_URI"
+
+popd

.evergreen/auth_aws/aws_tester.py

@@ -10,6 +10,7 @@
 from functools import partial
 
 from pymongo import MongoClient
+from pymongo.errors import OperationFailure
 from urllib.parse import quote_plus
 
 HERE = os.path.abspath(os.path.dirname(__file__))
@@ -50,7 +51,11 @@ def create_user(user, kwargs):
  print('Creating user', user)
  client = MongoClient(username="bob", password="pwd123")
  db = client['$external']
- db.command(dict(createUser=user, roles=[{"role": "read", "db": "aws"}]))
+ try:
+ db.command(dict(createUser=user, roles=[{"role": "read", "db": "aws"}]))
+ except OperationFailure as e:
+ if "already exists" not in e.details['errmsg']:
+ raise
  client.close()
 
  # Verify access.
@@ -65,7 +70,7 @@ def setup_assume_role():
  os.environ['AWS_SECRET_ACCESS_KEY'] = CONFIG[get_key("iam_auth_assume_aws_secret_access_key")]
 
  role_name = CONFIG[get_key("iam_auth_assume_role_name")]
- creds = _assume_role(role_name)
+ creds = _assume_role(role_name, quiet=True)
  with open(join(HERE, 'creds.json'), 'w') as fid:
  json.dump(creds, fid)
 
@@ -126,23 +131,25 @@ def setup_web_identity():
  print('ret was', ret)
  raise RuntimeError("Failed to unassign an instance profile from the current machine")
 
+ token_file = os.environ.get('AWS_WEB_IDENTITY_TOKEN_FILE', CONFIG[get_key('iam_web_identity_token_file')])
+
  # 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')]
+ AWS_WEB_IDENTITY_TOKEN_FILE=token_file
  )
 
  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_WEB_IDENTITY_TOKEN_FILE'] = token_file
  os.environ['AWS_ROLE_ARN'] = CONFIG[get_key("iam_auth_assume_web_role_name")]
 
- creds = _assume_role_with_web_identity()
+ creds = _assume_role_with_web_identity(True)
  with open(join(HERE, 'creds.json'), 'w') as fid:
  json.dump(creds, fid)
 

.evergreen/auth_aws/lib/aws_assume_role.py

@@ -13,21 +13,21 @@
 
 STS_DEFAULT_ROLE_NAME = "arn:aws:iam::579766882180:role/mark.benvenuto"
 
-def _assume_role(role_name):
+def _assume_role(role_name, quiet=False):
  sts_client = boto3.client("sts", region_name="us-east-1")
 
  response = sts_client.assume_role(RoleArn=role_name, RoleSessionName=str(uuid.uuid4()), DurationSeconds=900)
 
  creds = response["Credentials"]
  creds["Expiration"] = str(creds["Expiration"])
-
+ if quiet:
+ return creds
  print(f"""{{
  "AccessKeyId" : "{creds["AccessKeyId"]}",
  "SecretAccessKey" : "{creds["SecretAccessKey"]}",
  "SessionToken" : "{creds["SessionToken"]}",
  "Expiration" : "{creds["Expiration"]}"
 }}""")
- return creds
 
 
 def main() -> None:

.evergreen/auth_aws/lib/aws_assume_web_role.py

@@ -12,7 +12,7 @@
 
 LOGGER = logging.getLogger(__name__)
 
-def _assume_role_with_web_identity():
+def _assume_role_with_web_identity(quiet=False):
  sts_client = boto3.client("sts")
 
  token_file = os.environ['AWS_WEB_IDENTITY_TOKEN_FILE']
@@ -24,6 +24,8 @@ def _assume_role_with_web_identity():
 
  creds = response["Credentials"]
  creds["Expiration"] = str(creds["Expiration"])
+ if quiet:
+ return creds
 
  print(f"""{{
  "AccessKeyId" : "{creds["AccessKeyId"]}",

.evergreen/auth_aws/setup_secrets.sh

@@ -9,6 +9,7 @@ HERE=$(dirname $0)
 pushd $HERE
 . ./activate-authawsvenv.sh
 popd
+set -x
 echo "Getting secrets:" "$@"
 python $HERE/setup_secrets.py "$@"
 echo "Got secrets"

.gitignore

@@ -111,3 +111,4 @@ lb-expansion.yml
 orchestration.config
 atlas-expansion.yml
 secrets-export.sh
+token_file.txt