From 0c83cc85da04462ccc7252dca2b5d4faab69837f Mon Sep 17 00:00:00 2001
From: Pauline Ribeyre <4224001+paulineribeyre@users.noreply.github.com>
Date: Thu, 5 Dec 2024 10:28:45 -0600
Subject: [PATCH] assume role

---
 .secrets.baseline                |  4 ++--
 gen3workflow/config-default.yaml |  4 ++++
 gen3workflow/config.py           |  1 +
 gen3workflow/routes/s3.py        | 23 ++++++++++++++++-------
 4 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/.secrets.baseline b/.secrets.baseline
index 9b3abca..6064cdc 100644
--- a/.secrets.baseline
+++ b/.secrets.baseline
@@ -151,7 +151,7 @@
         "filename": "gen3workflow/config-default.yaml",
         "hashed_secret": "afc848c316af1a89d49826c5ae9d00ed769415f3",
         "is_verified": false,
-        "line_number": 27
+        "line_number": 31
       }
     ],
     "migrations/versions/e1886270d9d2_create_system_key_table.py": [
@@ -182,5 +182,5 @@
       }
     ]
   },
-  "generated_at": "2024-11-19T19:43:31Z"
+  "generated_at": "2024-12-05T16:27:30Z"
 }
diff --git a/gen3workflow/config-default.yaml b/gen3workflow/config-default.yaml
index 270384d..652c68b 100644
--- a/gen3workflow/config-default.yaml
+++ b/gen3workflow/config-default.yaml
@@ -16,6 +16,10 @@ MAX_IAM_KEYS_PER_USER: 2  # the default AWS AccessKeysPerUser quota is 2
 IAM_KEYS_LIFETIME_DAYS: 30
 USER_BUCKETS_REGION: us-east-1
 
+S3_ENDPOINTS_AWS_ROLE_ARN:
+S3_ENDPOINTS_AWS_ACCESS_KEY_ID:
+S3_ENDPOINTS_AWS_SECRET_ACCESS_KEY:
+
 #############
 # DATABASE #
 #############
diff --git a/gen3workflow/config.py b/gen3workflow/config.py
index fc747da..e722abc 100644
--- a/gen3workflow/config.py
+++ b/gen3workflow/config.py
@@ -48,6 +48,7 @@ def validate_top_level_configs(self):
                 "MAX_IAM_KEYS_PER_USER": {"type": "integer", "maximum": 100},
                 "IAM_KEYS_LIFETIME_DAYS": {"type": "integer"},
                 "USER_BUCKETS_REGION": {"type": "string"},
+                # TODO S3_ENDPOINTS_AWS_ROLE_ARN etc
                 "ARBORIST_URL": {"type": ["string", "null"]},
                 "TASK_IMAGE_WHITELIST": {"type": "array", "items": {"type": "string"}},
                 "TES_SERVER_URL": {"type": "string"},
diff --git a/gen3workflow/routes/s3.py b/gen3workflow/routes/s3.py
index b12fc73..c30af17 100644
--- a/gen3workflow/routes/s3.py
+++ b/gen3workflow/routes/s3.py
@@ -4,6 +4,7 @@
 import os
 import urllib.parse
 
+import boto3
 from fastapi import APIRouter, Request
 from fastapi.security import HTTPAuthorizationCredentials
 from botocore.credentials import Credentials
@@ -14,6 +15,7 @@
 
 from gen3workflow import aws_utils, logger
 from gen3workflow.auth import Auth
+from gen3workflow.config import config
 
 
 # TODO Generate a presigned URL if the request is a GET request, see https://cdis.slack.com/archives/D01DMJWKVB5/p1733169741227879 - is that required?
@@ -140,7 +142,7 @@ async def catch_all_v4(path: str, request: Request):
     #     headers['content-length'] = request.headers['content-length']
     # if 'x-amz-decoded-content-length' in request.headers:
     #     headers['x-amz-decoded-content-length'] = request.headers['x-amz-decoded-content-length']
-    
+
     # Ensure 'x-amz-date' is included in the headers (it's needed for signature calculation)
     amz_date = datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')
     headers['x-amz-date'] = amz_date
@@ -161,14 +163,21 @@ async def catch_all_v4(path: str, request: Request):
     # logger.debug(f"- Canonical Request:\n{canonical_request}")
 
     # AWS Credentials for signing
-    # TODO support either AWS IAM key or service account
-    credentials = Credentials(
-        access_key=os.environ.get('KEY'),
-        secret_key=os.environ.get('SECRET')
-    )
+    if config["S3_ENDPOINTS_AWS_ROLE_ARN"]:
+        sts_client = boto3.client('sts')
+        response = sts_client.assume_role(
+            RoleArn=config["S3_ENDPOINTS_AWS_ROLE_ARN"],
+            RoleSessionName='SessionName'
+        )
+        credentials = response['Credentials']
+    else:
+        credentials = Credentials(
+            access_key=config["S3_ENDPOINTS_AWS_ACCESS_KEY_ID"],
+            secret_key=config["S3_ENDPOINTS_AWS_SECRET_ACCESS_KEY"],
+        )
 
     # Create the string to sign based on the canonical request
-    region = 'us-east-1'
+    region = config["USER_BUCKETS_REGION"]
     service = 's3'
     date_stamp = headers['x-amz-date'][:8]  # The date portion (YYYYMMDD)
     string_to_sign = (