AWS Athena, added external id support (#1187)

opencybersecurityalliance · Nov 15, 2022 · c780d96 · c780d96
1 parent 119dd79
commit c780d96
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/stix_shifter_modules/aws_athena/configuration/config.json b/stix_shifter_modules/aws_athena/configuration/config.json
@@ -55,6 +55,10 @@
                 "type": "password",
                 "optional": true
             },
+            "aws_assume_role_external_id": {
+                "type": "password",
+                "optional": true
+            },
             "aws_secret_access_key": {
                 "type": "password"
             }

diff --git a/stix_shifter_modules/aws_athena/configuration/lang_en.json b/stix_shifter_modules/aws_athena/configuration/lang_en.json
@@ -55,6 +55,10 @@
                 "label": "AWS IAM Role",
                 "description": "AWS IAM Role is required for only AWS role-based authentication"
             },
+            "aws_assume_role_external_id": {
+                "label": "External ID for AWS Assume Role",
+                "description": "External ID is optional for using in AWS role-based authentication"
+            },
             "aws_secret_access_key": {
                 "label": "AWS Secret Access Key",
                 "description": "AWS Secret Access Key ID is required for both AWS key-based and role-based authentication"

diff --git a/stix_shifter_modules/aws_athena/stix_transmission/boto3_client.py b/stix_shifter_modules/aws_athena/stix_transmission/boto3_client.py
@@ -20,13 +20,21 @@ def __init__(self, connection, configuration):
                                       aws_secret_access_key=aws_secret_access_key,
                                       )
                 role_to_assume_arn = auth.get('aws_iam_role')
+                assume_role_external_id = auth.get('aws_assume_role_external_id')
                 prefix = 'AWS_'
                 letters = string.ascii_lowercase
                 role_session_name = prefix + ''.join(random.sample(letters, 4))
-                response = client.assume_role(
-                    RoleArn=role_to_assume_arn,
-                    RoleSessionName=role_session_name
-                )
+                if assume_role_external_id:
+                    response = client.assume_role(
+                        RoleArn=role_to_assume_arn,
+                        RoleSessionName=role_session_name,
+                        ExternalId=assume_role_external_id
+                    )
+                else:
+                    response = client.assume_role(
+                        RoleArn=role_to_assume_arn,
+                        RoleSessionName=role_session_name
+                    )
                 aws_creds = response['Credentials']
                 self.athena_client = boto3.client('athena',
                                                   aws_access_key_id=aws_creds['AccessKeyId'],