DRIVERS-3019 Apply auto formatting to python files (#536)

.evergreen/auth_aws/aws_tester.py

@@ -1,6 +1,7 @@
 """
 Script for testing MONGDOB-AWS authentication.
 """
+
 import argparse
 import json
 import os
@@ -14,11 +15,12 @@
 
 HERE = os.path.abspath(os.path.dirname(__file__))
 
+
 def join(*parts):
-    return os.path.join(*parts).replace(os.sep, '/')
+    return os.path.join(*parts).replace(os.sep, "/")
 
 
-sys.path.insert(0, join(HERE, 'lib'))
+sys.path.insert(0, join(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
@@ -32,7 +34,7 @@ def join(*parts):
 _USE_AWS_SECRETS = False
 
 try:
-    with open(join(HERE, 'aws_e2e_setup.json')) as fid:
+    with open(join(HERE, "aws_e2e_setup.json")) as fid:
         CONFIG = json.load(fid)
         get_key = partial(_get_key, uppercase=False)
 except FileNotFoundError:
@@ -48,36 +50,41 @@ def run(args, env):
 
 def create_user(user, kwargs):
     """Create a user and verify access."""
-    print('Creating user', user)
+    print("Creating user", user)
     client = MongoClient(username="bob", password="pwd123")
-    db = client['$external']
+    db = client["$external"]
     try:
         db.command(dict(createUser=user, roles=[{"role": "read", "db": "aws"}]))
     except OperationFailure as e:
-        if "already exists" not in e.details['errmsg']:
+        if "already exists" not in e.details["errmsg"]:
             raise
     client.close()
 
     # Verify access.
-    client = MongoClient(authMechanism='MONGODB-AWS', **kwargs)
-    client.aws.command('count', 'test')
+    client = MongoClient(authMechanism="MONGODB-AWS", **kwargs)
+    client.aws.command("count", "test")
     client.close()
 
 
 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")]
+    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, quiet=True)
-    with open(join(HERE, 'creds.json'), 'w') as fid:
+    with open(join(HERE, "creds.json"), "w") as fid:
         json.dump(creds, fid)
 
     # Create the user.
-    token = quote_plus(creds['SessionToken'])
-    kwargs = dict(username=creds["AccessKeyId"], password=creds["SecretAccessKey"],
-                  authmechanismproperties=f"AWS_SESSION_TOKEN:{token}")
+    token = quote_plus(creds["SessionToken"])
+    kwargs = dict(
+        username=creds["AccessKeyId"],
+        password=creds["SecretAccessKey"],
+        authmechanismproperties=f"AWS_SESSION_TOKEN:{token}",
+    )
     create_user(ASSUMED_ROLE, kwargs)
 
 
@@ -91,114 +98,131 @@ def setup_ec2():
 
 def setup_ecs():
     # Set up commands.
-    mongo_binaries = os.environ['MONGODB_BINARIES']
-    project_dir = os.environ['PROJECT_DIRECTORY']
+    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')]}"
 
     # Get the appropriate task definition based on the version of Ubuntu.
-    with open('/etc/lsb-release') as fid:
+    with open("/etc/lsb-release") as fid:
         text = fid.read()
-    if 'jammy' in text:
-        task_definition = CONFIG.get(get_key('iam_auth_ecs_task_definition_jammy'), None)
+    if "jammy" in text:
+        task_definition = CONFIG.get(
+            get_key("iam_auth_ecs_task_definition_jammy"), None
+        )
         if task_definition is None:
             raise ValueError('Please set "iam_auth_ecs_task_definition_jammy" variable')
-    elif 'focal' in text:
-        task_definition = CONFIG.get(get_key('iam_auth_ecs_task_definition_focal'), None)
+    elif "focal" in text:
+        task_definition = CONFIG.get(
+            get_key("iam_auth_ecs_task_definition_focal"), None
+        )
         # Fall back to previous task definition for backward compat.
         if task_definition is None:
-            task_definition = CONFIG[get_key('iam_auth_ecs_task_definition')]
+            task_definition = CONFIG[get_key("iam_auth_ecs_task_definition")]
     else:
-        raise ValueError('Unsupported ubuntu release')
+        raise ValueError("Unsupported ubuntu release")
     run_test_command = f"{base_command} -d -v run_e2e_test --cluster {CONFIG[get_key('iam_auth_ecs_cluster')]} --task_definition {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"
 
     # 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")],
+        AWS_SECRET_ACCESS_KEY=CONFIG[get_key("iam_auth_ecs_secret_access_key")],
+    )
 
     # Prune other containers
-    subprocess.check_call(['/bin/sh', '-c', run_prune_command], env=env)
+    subprocess.check_call(["/bin/sh", "-c", run_prune_command], env=env)
 
     # Run the test in a container
-    subprocess.check_call(['/bin/sh', '-c', run_test_command], env=env)
+    subprocess.check_call(["/bin/sh", "-c", run_test_command], env=env)
 
 
 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")]
+        password=CONFIG[get_key("iam_auth_ecs_secret_access_key")],
     )
     create_user(CONFIG[get_key("iam_auth_ecs_account_arn")], 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")])
-    ret = run(['lib/aws_unassign_instance_profile.py'], env)
+    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")
+        ],
+    )
+    ret = run(["lib/aws_unassign_instance_profile.py"], env)
     if ret == 2:
         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")
+        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')])
-    if os.name == "nt" and token_file.startswith('/tmp'):
+    token_file = os.environ.get(
+        "AWS_WEB_IDENTITY_TOKEN_FILE", CONFIG[get_key("iam_web_identity_token_file")]
+    )
+    if os.name == "nt" and token_file.startswith("/tmp"):
         token_file = token_file.replace("/tmp", "C:/cygwin/tmp/")
 
     # 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=token_file
+        AWS_WEB_IDENTITY_TOKEN_FILE=token_file,
     )
 
-    ret = run(['lib/aws_handle_oidc_creds.py', 'token'], env)
+    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'] = token_file
-    os.environ['AWS_ROLE_ARN'] = CONFIG[get_key("iam_auth_assume_web_role_name")]
+    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(True)
-    with open(join(HERE, 'creds.json'), 'w') as fid:
+    with open(join(HERE, "creds.json"), "w") as fid:
         json.dump(creds, fid)
 
     # Create the user.
-    token = quote_plus(creds['SessionToken'])
-    kwargs = dict(username=creds["AccessKeyId"], password=creds["SecretAccessKey"],
-                  authmechanismproperties=f"AWS_SESSION_TOKEN:{token}")
+    token = quote_plus(creds["SessionToken"])
+    kwargs = dict(
+        username=creds["AccessKeyId"],
+        password=creds["SecretAccessKey"],
+        authmechanismproperties=f"AWS_SESSION_TOKEN:{token}",
+    )
     create_user(ASSUMED_WEB_ROLE, kwargs)
 
 
 def main():
-    parser = argparse.ArgumentParser(description='MONGODB-AWS tester.')
+    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", help="Assume role test")
     run_assume_role_cmd.set_defaults(func=setup_assume_role)
 
-    run_ec2_cmd = sub.add_parser('ec2', help='EC2 test')
+    run_ec2_cmd = sub.add_parser("ec2", help="EC2 test")
     run_ec2_cmd.set_defaults(func=setup_ec2)
 
-    run_ecs_cmd = sub.add_parser('ecs', help='ECS test')
+    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", 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 = 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()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

.evergreen/auth_aws/lib/aws_assign_instance_profile.py

@@ -15,14 +15,18 @@
 import botocore
 from util import get_key as _get_key
 
-sys.path.insert(1, os.path.join(sys.path[0], '..'))
+sys.path.insert(1, os.path.join(sys.path[0], ".."))
 
 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()
+    return (
+        urllib.request.urlopen("http://169.254.169.254/latest/meta-data/instance-id")
+        .read()
+        .decode()
+    )
 
 
 def _has_instance_profile():
@@ -62,7 +66,7 @@ def _wait_instance_profile():
 
 def _handle_config():
     try:
-        with open(os.path.join(HERE, '..', 'aws_e2e_setup.json')) as fid:
+        with open(os.path.join(HERE, "..", "aws_e2e_setup.json")) as fid:
             CONFIG = json.load(fid)
             get_key = partial(_get_key, uppercase=False)
 
@@ -71,34 +75,41 @@ def _handle_config():
         get_key = partial(_get_key, uppercase=True)
 
     try:
-        os.environ.setdefault('AWS_ACCESS_KEY_ID', CONFIG[get_key('iam_auth_ec2_instance_account')])
-        os.environ.setdefault('AWS_SECRET_ACCESS_KEY',
-                              CONFIG[get_key('iam_auth_ec2_instance_secret_access_key')])
-        return CONFIG[get_key('iam_auth_ec2_instance_profile')]
+        os.environ.setdefault(
+            "AWS_ACCESS_KEY_ID", CONFIG[get_key("iam_auth_ec2_instance_account")]
+        )
+        os.environ.setdefault(
+            "AWS_SECRET_ACCESS_KEY",
+            CONFIG[get_key("iam_auth_ec2_instance_secret_access_key")],
+        )
+        return CONFIG[get_key("iam_auth_ec2_instance_profile")]
     except Exception as e:
         print(e)
-        return ''
+        return ""
 
 
 DEFAULT_ARN = _handle_config()
 
 
 def _assign_instance_policy(iam_instance_arn=DEFAULT_ARN):
     if _has_instance_profile():
-        print("IMPORTANT: Found machine already has instance profile, skipping the assignment")
+        print(
+            "IMPORTANT: Found machine already has instance profile, skipping the assignment"
+        )
         return
 
     instance_id = _get_local_instance_id()
 
-    ec2_client = boto3.client("ec2", 'us-east-1')
+    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
     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)
 
@@ -115,12 +126,21 @@ def _assign_instance_policy(iam_instance_arn=DEFAULT_ARN):
 def main() -> None:
     """Execute Main entry point."""
 
-    parser = argparse.ArgumentParser(description='IAM Assign Instance frontend.')
-
-    parser.add_argument('-v', "--verbose", action='store_true', help="Enable verbose logging")
-    parser.add_argument('-d', "--debug", action='store_true', help="Enable debug logging")
-
-    parser.add_argument('--instance_profile_arn', type=str, help="Name of instance profile", default=DEFAULT_ARN)
+    parser = argparse.ArgumentParser(description="IAM Assign Instance frontend.")
+
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose logging"
+    )
+    parser.add_argument(
+        "-d", "--debug", action="store_true", help="Enable debug logging"
+    )
+
+    parser.add_argument(
+        "--instance_profile_arn",
+        type=str,
+        help="Name of instance profile",
+        default=DEFAULT_ARN,
+    )
 
     args = parser.parse_args()
 

.evergreen/auth_aws/lib/aws_assume_role.py

@@ -12,10 +12,13 @@
 
 STS_DEFAULT_ROLE_NAME = "arn:aws:iam::579766882180:role/mark.benvenuto"
 
+
 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)
+    response = sts_client.assume_role(
+        RoleArn=role_name, RoleSessionName=str(uuid.uuid4()), DurationSeconds=900
+    )
 
     creds = response["Credentials"]
     creds["Expiration"] = str(creds["Expiration"])
@@ -32,12 +35,18 @@ def _assume_role(role_name, quiet=False):
 def main() -> None:
     """Execute Main entry point."""
 
-    parser = argparse.ArgumentParser(description='Assume Role frontend.')
+    parser = argparse.ArgumentParser(description="Assume Role frontend.")
 
-    parser.add_argument('-v', "--verbose", action='store_true', help="Enable verbose logging")
-    parser.add_argument('-d', "--debug", action='store_true', help="Enable debug logging")
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose logging"
+    )
+    parser.add_argument(
+        "-d", "--debug", action="store_true", help="Enable debug logging"
+    )
 
-    parser.add_argument('--role_name', type=str, default=STS_DEFAULT_ROLE_NAME, help="Role to assume")
+    parser.add_argument(
+        "--role_name", type=str, default=STS_DEFAULT_ROLE_NAME, help="Role to assume"
+    )
 
     args = parser.parse_args()