From 28fb39b1441e5ab152e5de945981d7e9a34a18f2 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Tue, 15 Oct 2024 14:37:03 +1100
Subject: [PATCH 01/11] Add time cache decorator

---
 packages/opal-common/opal_common/utils.py | 28 ++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/packages/opal-common/opal_common/utils.py b/packages/opal-common/opal_common/utils.py
index 3897c058f..3c843ce44 100644
--- a/packages/opal-common/opal_common/utils.py
+++ b/packages/opal-common/opal_common/utils.py
@@ -8,7 +8,9 @@
 import threading
 from datetime import datetime
 from hashlib import sha1
-from typing import Coroutine, Dict, List, Tuple
+from typing import Callable, Coroutine, Dict, List, Tuple
+import functools
+import time
 
 import aiohttp
 
@@ -275,3 +277,27 @@ def run_coro(self, coro: Coroutine):
         run_coro() is thread-safe.
         """
         return asyncio.run_coroutine_threadsafe(coro, loop=self.loop).result()
+
+
+def time_cache(ttl: float):
+    """
+    This decorator a wrapper around lru_cache that makes it time sensitive.
+
+    ttl is in seconds
+    """
+
+    def inner(func: Callable):
+        # instead of directly caching the function, a time "hash" is
+        # also passed in as a param that will invalidate the cache
+        # after at most ttl seconds
+        @functools.lru_cache
+        def wrapped(*args, __ttl_hash=None, **kwargs):
+            return func(*args, **kwargs)
+
+        def ret(*args, **kwargs):
+            ttl_hash = round(time.time() / ttl)
+            return wrapped(*args, **kwargs, __ttl_hash=ttl_hash)
+
+        return ret
+
+    return inner

From 3350181d3efd22cd07362a1993885363f8f37ac9 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 16 Oct 2024 14:24:32 +1100
Subject: [PATCH 02/11] Add extra config required for iam auth

---
 .../opal_common/sources/api_policy_source.py           |  4 ++++
 packages/opal-server/opal_server/config.py             | 10 ++++++++++
 .../opal-server/opal_server/policy/watcher/factory.py  |  2 ++
 3 files changed, 16 insertions(+)

diff --git a/packages/opal-common/opal_common/sources/api_policy_source.py b/packages/opal-common/opal_common/sources/api_policy_source.py
index 7adc9ad70..3418eb3df 100644
--- a/packages/opal-common/opal_common/sources/api_policy_source.py
+++ b/packages/opal-common/opal_common/sources/api_policy_source.py
@@ -53,6 +53,8 @@ def __init__(
         token: Optional[str] = None,
         token_id: Optional[str] = None,
         region: Optional[str] = None,
+        role_arn: Optional[str] = None,
+        token_file: Optional[str] = None,
         bundle_server_type: Optional[PolicyBundleServerType] = None,
         policy_bundle_path=".",
         policy_bundle_git_add_pattern="*",
@@ -66,6 +68,8 @@ def __init__(
         self.token_id = token_id
         self.server_type = bundle_server_type
         self.region = region
+        self.role_arn = role_arn
+        self.token_file = token_file
         self.bundle_hash = None
         self.etag = None
         self.tmp_bundle_path = Path(policy_bundle_path)
diff --git a/packages/opal-server/opal_server/config.py b/packages/opal-server/opal_server/config.py
index b272915ad..49a3494a4 100644
--- a/packages/opal-server/opal_server/config.py
+++ b/packages/opal-server/opal_server/config.py
@@ -133,6 +133,16 @@ class OpalServerConfig(Confi):
         "us-east-1",
         description="The AWS region of the S3 bucket",
     )
+    POLICY_BUNDLE_AWS_ROLE_ARN = confi.str(
+        "AWS_ROLE_ARN",
+        None,
+        description="The IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
+    )
+    POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE = confi.str(
+        "AWS_WEB_IDENTITY_TOKEN_FILE",
+        None,
+        description="The oidc token for the IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
+    )
     POLICY_BUNDLE_TMP_PATH = confi.str(
         "POLICY_BUNDLE_TMP_PATH",
         "/tmp/bundle.tar.gz",
diff --git a/packages/opal-server/opal_server/policy/watcher/factory.py b/packages/opal-server/opal_server/policy/watcher/factory.py
index 6d94d6fc4..3ce36cfd7 100644
--- a/packages/opal-server/opal_server/policy/watcher/factory.py
+++ b/packages/opal-server/opal_server/policy/watcher/factory.py
@@ -129,6 +129,8 @@ def setup_watcher_task(
             policy_bundle_path=opal_server_config.POLICY_BUNDLE_TMP_PATH,
             policy_bundle_git_add_pattern=opal_server_config.POLICY_BUNDLE_GIT_ADD_PATTERN,
             region=policy_bundle_aws_region,
+            role_arn=opal_server_config.POLICY_BUNDLE_AWS_ROLE_ARN,
+            token_file=opal_server_config.POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE,,
         )
     else:
         raise ValueError("Unknown value for OPAL_POLICY_SOURCE_TYPE")

From 7e4cd88b6fa7e508299c812098b43c8acbd5fc8c Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 16 Oct 2024 14:25:28 +1100
Subject: [PATCH 03/11] Make build_auth_headers async

---
 packages/opal-common/opal_common/sources/api_policy_source.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/opal-common/opal_common/sources/api_policy_source.py b/packages/opal-common/opal_common/sources/api_policy_source.py
index 3418eb3df..738e50645 100644
--- a/packages/opal-common/opal_common/sources/api_policy_source.py
+++ b/packages/opal-common/opal_common/sources/api_policy_source.py
@@ -130,7 +130,7 @@ async def api_update_policy(self) -> Tuple[bool, str, str]:
                     )
                     raise
 
-    def build_auth_headers(self, token=None, path=None):
+    async def build_auth_headers(self, token=None, path=None):
         # if it's a simple HTTP server with a bearer token
         if self.server_type == PolicyBundleServerType.HTTP and token is not None:
             return tuple_to_dict(get_authorization_header(token))
@@ -170,7 +170,7 @@ async def fetch_policy_bundle_from_api_source(
         """
         path = "bundle.tar.gz"
 
-        auth_headers = self.build_auth_headers(token=token, path=path)
+        auth_headers = await self.build_auth_headers(token=token, path=path)
         etag_headers = (
             {"ETag": self.etag, "If-None-Match": self.etag} if self.etag else {}
         )

From bad572e1fa068c98faba36a272281141c0090765 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 16 Oct 2024 14:27:35 +1100
Subject: [PATCH 04/11] Add method to get temp credentials via iam auth

---
 .../opal_common/sources/api_policy_source.py  | 70 +++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/packages/opal-common/opal_common/sources/api_policy_source.py b/packages/opal-common/opal_common/sources/api_policy_source.py
index 738e50645..58f3c7670 100644
--- a/packages/opal-common/opal_common/sources/api_policy_source.py
+++ b/packages/opal-common/opal_common/sources/api_policy_source.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import Optional, Tuple
 from urllib.parse import urlparse
+from xml.etree import ElementTree
 
 import aiohttp
 from fastapi import status
@@ -130,6 +131,62 @@ async def api_update_policy(self) -> Tuple[bool, str, str]:
                     )
                     raise
 
+    async def get_temporary_sts_credentials(self) -> Tuple[str, str]:
+        assert self.token_file
+        assert self.role_arn
+        assert self.region
+
+        with open(self.token_file) as token_file:
+            token = token_file.read()
+
+        sts_url = f"sts.{self.region}.amazonaws.com"
+        params = (
+            {
+                "Action": "AssumeRoleWithWebIdentity",
+                "DurationSeconds": 3600,
+                "RoleSessionName": "Opal",
+                "RoleArn": self.role_arn,
+                "WebIdentityToken": token,
+                "Version": "2011-06-15",
+            },
+        )
+
+        async with aiohttp.ClientSession() as session:
+            try:
+                async with session.get(
+                    f"{sts_url}",
+                    params=params,
+                ) as response:
+                    if response.status == status.HTTP_404_NOT_FOUND:
+                        logger.warning(
+                            "requested url not found: {sts_url}",
+                            sts_url=sts_url,
+                        )
+                        raise HTTPException(
+                            status_code=status.HTTP_404_NOT_FOUND,
+                            detail=f"requested url not found: {sts_url}",
+                        )
+
+                    body = await response.read()
+
+                    et = ElementTree.parse(body)
+                    credentials = et.find(
+                        "/AssumeRoleWithWebIdentityResponse/AssumeRoleWithWebIdentityResult/Credentials"
+                    )
+                    assert credentials
+
+                    assert (id := credentials.findtext("AccessKeyId"))
+                    assert (key := credentials.findtext("SecretAccessKey"))
+
+            except (aiohttp.ClientError, HTTPException) as e:
+                logger.warning("server connection error: {err}", err=repr(e))
+                raise
+            except Exception as e:
+                logger.error("unexpected server connection error: {err}", err=repr(e))
+                raise
+
+        return id, key
+
     async def build_auth_headers(self, token=None, path=None):
         # if it's a simple HTTP server with a bearer token
         if self.server_type == PolicyBundleServerType.HTTP and token is not None:
@@ -147,6 +204,19 @@ async def build_auth_headers(self, token=None, path=None):
             return build_aws_rest_auth_headers(
                 self.token_id, token, host, path, self.region
             )
+        elif (
+            self.server_type == PolicyBundleServerType.AWS_S3
+            and self.role_arn is not None
+            and self.token_file is not None
+            and self.region is not None
+        ):
+            split_url = urlparse(self.remote_source_url)
+            host = split_url.netloc
+            path = split_url.path + "/" + path
+
+            id, key = await self.get_temporary_sts_credentials()
+
+            return build_aws_rest_auth_headers(id, key, host, path, self.region)
         else:
             return {}
 

From 1278ca86595bfce80a9edda8b69a3af11b1ee79a Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 16 Oct 2024 14:37:00 +1100
Subject: [PATCH 05/11] remove prefix from aws provided env vars

---
 packages/opal-server/opal_server/config.py                 | 5 +++--
 packages/opal-server/opal_server/policy/watcher/factory.py | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/packages/opal-server/opal_server/config.py b/packages/opal-server/opal_server/config.py
index 49a3494a4..c1b982260 100644
--- a/packages/opal-server/opal_server/config.py
+++ b/packages/opal-server/opal_server/config.py
@@ -8,6 +8,7 @@
 from opal_common.schemas.webhook import GitWebhookRequestParams
 
 confi = Confi(prefix="OPAL_")
+confi_no_prefix = Confi(prefix=None)
 
 
 class PolicySourceTypes(str, Enum):
@@ -133,12 +134,12 @@ class OpalServerConfig(Confi):
         "us-east-1",
         description="The AWS region of the S3 bucket",
     )
-    POLICY_BUNDLE_AWS_ROLE_ARN = confi.str(
+    POLICY_BUNDLE_AWS_ROLE_ARN = confi_no_prefix.str(
         "AWS_ROLE_ARN",
         None,
         description="The IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
     )
-    POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE = confi.str(
+    POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE = confi_no_prefix.str(
         "AWS_WEB_IDENTITY_TOKEN_FILE",
         None,
         description="The oidc token for the IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
diff --git a/packages/opal-server/opal_server/policy/watcher/factory.py b/packages/opal-server/opal_server/policy/watcher/factory.py
index 3ce36cfd7..14688c006 100644
--- a/packages/opal-server/opal_server/policy/watcher/factory.py
+++ b/packages/opal-server/opal_server/policy/watcher/factory.py
@@ -130,7 +130,7 @@ def setup_watcher_task(
             policy_bundle_git_add_pattern=opal_server_config.POLICY_BUNDLE_GIT_ADD_PATTERN,
             region=policy_bundle_aws_region,
             role_arn=opal_server_config.POLICY_BUNDLE_AWS_ROLE_ARN,
-            token_file=opal_server_config.POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE,,
+            token_file=opal_server_config.POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE,
         )
     else:
         raise ValueError("Unknown value for OPAL_POLICY_SOURCE_TYPE")

From c11b03cd565d1f39db72dbc1506a080c69282a56 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Mon, 21 Oct 2024 11:26:01 +1100
Subject: [PATCH 06/11] more logs, and set content type

---
 .../opal-common/opal_common/sources/api_policy_source.py  | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/packages/opal-common/opal_common/sources/api_policy_source.py b/packages/opal-common/opal_common/sources/api_policy_source.py
index 58f3c7670..38bdd7da7 100644
--- a/packages/opal-common/opal_common/sources/api_policy_source.py
+++ b/packages/opal-common/opal_common/sources/api_policy_source.py
@@ -156,6 +156,7 @@ async def get_temporary_sts_credentials(self) -> Tuple[str, str]:
                 async with session.get(
                     f"{sts_url}",
                     params=params,
+                    headers={"Content-Type": "application/xml"},
                 ) as response:
                     if response.status == status.HTTP_404_NOT_FOUND:
                         logger.warning(
@@ -185,6 +186,7 @@ async def get_temporary_sts_credentials(self) -> Tuple[str, str]:
                 logger.error("unexpected server connection error: {err}", err=repr(e))
                 raise
 
+        logger.info("Successfully generated temporary AWS credentials")
         return id, key
 
     async def build_auth_headers(self, token=None, path=None):
@@ -197,6 +199,8 @@ async def build_auth_headers(self, token=None, path=None):
             and token is not None
             and self.token_id is not None
         ):
+            logger.info("Using provided token to login to AWS_S3")
+
             split_url = urlparse(self.remote_source_url)
             host = split_url.netloc
             path = split_url.path + "/" + path
@@ -208,8 +212,9 @@ async def build_auth_headers(self, token=None, path=None):
             self.server_type == PolicyBundleServerType.AWS_S3
             and self.role_arn is not None
             and self.token_file is not None
-            and self.region is not None
         ):
+            logger.info("Using IAM Web auth to login to AWS_S3")
+
             split_url = urlparse(self.remote_source_url)
             host = split_url.netloc
             path = split_url.path + "/" + path
@@ -218,6 +223,7 @@ async def build_auth_headers(self, token=None, path=None):
 
             return build_aws_rest_auth_headers(id, key, host, path, self.region)
         else:
+            logger.info("Not authenticating on bundle endpoint")
             return {}
 
     async def fetch_policy_bundle_from_api_source(

From cdf669fbbe354cf9e510cb0507985352de5a6937 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 30 Oct 2024 11:48:28 +1100
Subject: [PATCH 07/11] Force direct read of env vars

I can't seem to figure out how non-prefixed vars should be read from
when using confi.
---
 packages/opal-server/opal_server/config.py    | 27 ++++++++++---------
 .../opal_server/policy/watcher/factory.py     |  9 ++++---
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/packages/opal-server/opal_server/config.py b/packages/opal-server/opal_server/config.py
index c1b982260..62cb91ada 100644
--- a/packages/opal-server/opal_server/config.py
+++ b/packages/opal-server/opal_server/config.py
@@ -8,7 +8,7 @@
 from opal_common.schemas.webhook import GitWebhookRequestParams
 
 confi = Confi(prefix="OPAL_")
-confi_no_prefix = Confi(prefix=None)
+# confi_no_prefix = Confi(prefix="")
 
 
 class PolicySourceTypes(str, Enum):
@@ -52,7 +52,8 @@ class OpalServerConfig(Confi):
     AUTH_PRIVATE_KEY_PASSPHRASE = confi.str("AUTH_PRIVATE_KEY_PASSPHRASE", None)
 
     AUTH_PRIVATE_KEY = confi.delay(
-        lambda AUTH_PRIVATE_KEY_FORMAT=None, AUTH_PRIVATE_KEY_PASSPHRASE="": confi.private_key(
+        lambda AUTH_PRIVATE_KEY_FORMAT=None,
+        AUTH_PRIVATE_KEY_PASSPHRASE="": confi.private_key(
             "AUTH_PRIVATE_KEY",
             default=None,
             key_format=AUTH_PRIVATE_KEY_FORMAT,
@@ -134,16 +135,16 @@ class OpalServerConfig(Confi):
         "us-east-1",
         description="The AWS region of the S3 bucket",
     )
-    POLICY_BUNDLE_AWS_ROLE_ARN = confi_no_prefix.str(
-        "AWS_ROLE_ARN",
-        None,
-        description="The IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
-    )
-    POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE = confi_no_prefix.str(
-        "AWS_WEB_IDENTITY_TOKEN_FILE",
-        None,
-        description="The oidc token for the IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
-    )
+    # POLICY_BUNDLE_AWS_ROLE_ARN = confi_no_prefix.str(
+    #     "AWS_ROLE_ARN",
+    #     None,
+    #     description="The IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
+    # )
+    # POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE = confi_no_prefix.str(
+    #     "AWS_WEB_IDENTITY_TOKEN_FILE",
+    #     None,
+    #     description="The oidc token for the IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
+    # )
     POLICY_BUNDLE_TMP_PATH = confi.str(
         "POLICY_BUNDLE_TMP_PATH",
         "/tmp/bundle.tar.gz",
@@ -328,4 +329,4 @@ def on_load(self):
             self.SERVER_BIND_PORT = int(self.SERVER_PORT)
 
 
-opal_server_config = OpalServerConfig(prefix="OPAL_")
+opal_server_config = OpalServerConfig(prefix="OPAL_")
\ No newline at end of file
diff --git a/packages/opal-server/opal_server/policy/watcher/factory.py b/packages/opal-server/opal_server/policy/watcher/factory.py
index 14688c006..7b3af1ae4 100644
--- a/packages/opal-server/opal_server/policy/watcher/factory.py
+++ b/packages/opal-server/opal_server/policy/watcher/factory.py
@@ -1,5 +1,6 @@
 from functools import partial
 from typing import Any, List, Optional
+import os
 
 from fastapi_websocket_pubsub.pub_sub_server import PubSubEndpoint
 from opal_common.confi.confi import load_conf_if_none
@@ -129,8 +130,10 @@ def setup_watcher_task(
             policy_bundle_path=opal_server_config.POLICY_BUNDLE_TMP_PATH,
             policy_bundle_git_add_pattern=opal_server_config.POLICY_BUNDLE_GIT_ADD_PATTERN,
             region=policy_bundle_aws_region,
-            role_arn=opal_server_config.POLICY_BUNDLE_AWS_ROLE_ARN,
-            token_file=opal_server_config.POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE,
+            # role_arn=opal_server_config.POLICY_BUNDLE_AWS_ROLE_ARN,
+            # token_file=opal_server_config.POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE,
+            role_arn=os.getenv("AWS_ROLE_ARN"),
+            token_file=os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE"),
         )
     else:
         raise ValueError("Unknown value for OPAL_POLICY_SOURCE_TYPE")
@@ -142,4 +145,4 @@ def setup_watcher_task(
             bundle_ignore=bundle_ignore,
         )
     )
-    return PolicyWatcherTask(watcher, pubsub_endpoint)
+    return PolicyWatcherTask(watcher, pubsub_endpoint)
\ No newline at end of file

From be759326ea8f1c20cb1b437e160e00cb5c5bef18 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 30 Oct 2024 15:36:59 +1100
Subject: [PATCH 08/11] Change ttl cache to work for async

---
 .../opal-common/opal_common/sources/api_policy_source.py | 2 ++
 packages/opal-common/opal_common/utils.py                | 9 +++++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/packages/opal-common/opal_common/sources/api_policy_source.py b/packages/opal-common/opal_common/sources/api_policy_source.py
index 38bdd7da7..0bec8eee6 100644
--- a/packages/opal-common/opal_common/sources/api_policy_source.py
+++ b/packages/opal-common/opal_common/sources/api_policy_source.py
@@ -18,6 +18,7 @@
     hash_file,
     throw_if_bad_status_code,
     tuple_to_dict,
+    async_time_cache,
 )
 from opal_server.config import PolicyBundleServerType
 from tenacity import AsyncRetrying
@@ -132,6 +133,7 @@ async def api_update_policy(self) -> Tuple[bool, str, str]:
                     raise
 
     async def get_temporary_sts_credentials(self) -> Tuple[str, str]:
+    @async_time_cache(ttl=3000)
         assert self.token_file
         assert self.role_arn
         assert self.region
diff --git a/packages/opal-common/opal_common/utils.py b/packages/opal-common/opal_common/utils.py
index 3c843ce44..4733dabb4 100644
--- a/packages/opal-common/opal_common/utils.py
+++ b/packages/opal-common/opal_common/utils.py
@@ -279,20 +279,21 @@ def run_coro(self, coro: Coroutine):
         return asyncio.run_coroutine_threadsafe(coro, loop=self.loop).result()
 
 
-def time_cache(ttl: float):
+def async_time_cache(ttl: float):
     """
     This decorator a wrapper around lru_cache that makes it time sensitive.
 
     ttl is in seconds
     """
 
-    def inner(func: Callable):
+    def decorator(func: Callable):
         # instead of directly caching the function, a time "hash" is
         # also passed in as a param that will invalidate the cache
         # after at most ttl seconds
         @functools.lru_cache
         def wrapped(*args, __ttl_hash=None, **kwargs):
-            return func(*args, **kwargs)
+            coro = func(*args, **kwargs)
+            return asyncio.ensure_future(coro)
 
         def ret(*args, **kwargs):
             ttl_hash = round(time.time() / ttl)
@@ -300,4 +301,4 @@ def ret(*args, **kwargs):
 
         return ret
 
-    return inner
+    return decorator
\ No newline at end of file

From c7dedb7d8aeaefad4eb33c5bba091b98eb078b81 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 30 Oct 2024 15:41:33 +1100
Subject: [PATCH 09/11] Allow aws header builder to take session token

---
 packages/opal-common/opal_common/utils.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/packages/opal-common/opal_common/utils.py b/packages/opal-common/opal_common/utils.py
index 4733dabb4..4f9e1f56e 100644
--- a/packages/opal-common/opal_common/utils.py
+++ b/packages/opal-common/opal_common/utils.py
@@ -59,7 +59,12 @@ def get_authorization_header(token: str) -> Tuple[str, str]:
 
 
 def build_aws_rest_auth_headers(
-    key_id: str, secret_key: str, host: str, path: str, region: str
+    key_id: str,
+    secret_key: str,
+    host: str,
+    path: str,
+    region: str,
+    token: str | None,
 ):
     """Use the AWS signature algorithm (https://docs.aws.amazon.com/AmazonS3/la
     test/userguide/RESTAuthentication.html) to generate the hTTP headers.
@@ -69,6 +74,7 @@ def build_aws_rest_auth_headers(
         secret_key (str): Secret key (aka password) of an account in the S3 service.
         host (str): S3 storage host
         path (str): path to bundle file in s3 storage (including bucket)
+        token (str | None): Optional session token when using temporary credential.
 
     Returns: http headers
     """
@@ -93,6 +99,10 @@ def getSignatureKey(key, dateStamp, regionName, serviceName):
     canonical_headers = "host:" + host + "\n" + "x-amz-date:" + amzdate + "\n"
     signed_headers = "host;x-amz-date"
 
+    if token:
+        canonical_headers += f"x-amz-security-token:{token}\n"
+        signed_headers += ";x-amz-security-token"
+
     payload_hash = hashlib.sha256("".encode("utf-8")).hexdigest()
 
     canonical_request = (
@@ -140,8 +150,13 @@ def getSignatureKey(key, dateStamp, regionName, serviceName):
         + signature
     )
 
+    token_header: dict[str, str] = {}
+    if token:
+        token_header["x-amz-security-token"] = token
+
     return {
         "x-amz-date": amzdate,
+        **token_header,
         "x-amz-content-sha256": SHA256_EMPTY,
         "Authorization": authorization_header,
     }

From 160ba209718897d6f5d7265c5984d8b7e6fa34f7 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 30 Oct 2024 15:44:09 +1100
Subject: [PATCH 10/11] Fix XML, async file handling, use session token

---
 .../opal_common/sources/api_policy_source.py  | 52 +++++++++++--------
 1 file changed, 31 insertions(+), 21 deletions(-)

diff --git a/packages/opal-common/opal_common/sources/api_policy_source.py b/packages/opal-common/opal_common/sources/api_policy_source.py
index 0bec8eee6..e9c1b8765 100644
--- a/packages/opal-common/opal_common/sources/api_policy_source.py
+++ b/packages/opal-common/opal_common/sources/api_policy_source.py
@@ -5,6 +5,7 @@
 from xml.etree import ElementTree
 
 import aiohttp
+import aiofiles
 from fastapi import status
 from fastapi.exceptions import HTTPException
 from opal_common.git_utils.tar_file_to_local_git_extractor import (
@@ -132,31 +133,29 @@ async def api_update_policy(self) -> Tuple[bool, str, str]:
                     )
                     raise
 
-    async def get_temporary_sts_credentials(self) -> Tuple[str, str]:
     @async_time_cache(ttl=3000)
+    async def get_temporary_sts_credentials(self) -> tuple[str, str, str]:
         assert self.token_file
         assert self.role_arn
         assert self.region
 
-        with open(self.token_file) as token_file:
-            token = token_file.read()
+        async with aiofiles.open(self.token_file) as token_file:
+            token = await token_file.read()
 
         sts_url = f"sts.{self.region}.amazonaws.com"
-        params = (
-            {
-                "Action": "AssumeRoleWithWebIdentity",
-                "DurationSeconds": 3600,
-                "RoleSessionName": "Opal",
-                "RoleArn": self.role_arn,
-                "WebIdentityToken": token,
-                "Version": "2011-06-15",
-            },
-        )
+        params: dict[str, str] = {
+            "Action": "AssumeRoleWithWebIdentity",
+            "DurationSeconds": "3600",
+            "RoleSessionName": "Opal",
+            "RoleArn": self.role_arn,
+            "WebIdentityToken": token,
+            "Version": "2011-06-15",
+        }
 
         async with aiohttp.ClientSession() as session:
             try:
                 async with session.get(
-                    f"{sts_url}",
+                    f"https://{sts_url}",
                     params=params,
                     headers={"Content-Type": "application/xml"},
                 ) as response:
@@ -172,14 +171,22 @@ async def get_temporary_sts_credentials(self) -> Tuple[str, str]:
 
                     body = await response.read()
 
-                    et = ElementTree.parse(body)
+                    # the default aws xml namespace
+                    ns = {"": "https://sts.amazonaws.com/doc/2011-06-15/"}
+
+                    et = ElementTree.fromstring(body)
                     credentials = et.find(
-                        "/AssumeRoleWithWebIdentityResponse/AssumeRoleWithWebIdentityResult/Credentials"
+                        "AssumeRoleWithWebIdentityResult/Credentials", ns
                     )
                     assert credentials
 
-                    assert (id := credentials.findtext("AccessKeyId"))
-                    assert (key := credentials.findtext("SecretAccessKey"))
+                    id = credentials.findtext("AccessKeyId", namespaces=ns)
+                    key = credentials.findtext("SecretAccessKey", namespaces=ns)
+                    session_token = credentials.findtext("SessionToken", namespaces=ns)
+
+                    assert id
+                    assert key
+                    assert session_token
 
             except (aiohttp.ClientError, HTTPException) as e:
                 logger.warning("server connection error: {err}", err=repr(e))
@@ -189,7 +196,7 @@ async def get_temporary_sts_credentials(self) -> Tuple[str, str]:
                 raise
 
         logger.info("Successfully generated temporary AWS credentials")
-        return id, key
+        return id, key, session_token
 
     async def build_auth_headers(self, token=None, path=None):
         # if it's a simple HTTP server with a bearer token
@@ -214,6 +221,7 @@ async def build_auth_headers(self, token=None, path=None):
             self.server_type == PolicyBundleServerType.AWS_S3
             and self.role_arn is not None
             and self.token_file is not None
+            and self.region is not None
         ):
             logger.info("Using IAM Web auth to login to AWS_S3")
 
@@ -221,9 +229,11 @@ async def build_auth_headers(self, token=None, path=None):
             host = split_url.netloc
             path = split_url.path + "/" + path
 
-            id, key = await self.get_temporary_sts_credentials()
+            id, key, session_token = await self.get_temporary_sts_credentials()
 
-            return build_aws_rest_auth_headers(id, key, host, path, self.region)
+            return build_aws_rest_auth_headers(
+                id, key, host, path, self.region, session_token
+            )
         else:
             logger.info("Not authenticating on bundle endpoint")
             return {}

From d6b71c1e790a0bede442edae013f6e010a184db8 Mon Sep 17 00:00:00 2001
From: "ibrahim.fuad" <ibrahim.fuad@saberastro.com>
Date: Wed, 30 Oct 2024 17:36:14 +1100
Subject: [PATCH 11/11] Docs, typos and naming.

---
 .../opal_common/sources/api_policy_source.py  | 41 +++++++++++++------
 packages/opal-common/opal_common/utils.py     |  4 +-
 packages/opal-server/opal_server/config.py    | 25 +++++------
 .../opal_server/policy/watcher/factory.py     |  8 ++--
 4 files changed, 46 insertions(+), 32 deletions(-)

diff --git a/packages/opal-common/opal_common/sources/api_policy_source.py b/packages/opal-common/opal_common/sources/api_policy_source.py
index e9c1b8765..22d2c06da 100644
--- a/packages/opal-common/opal_common/sources/api_policy_source.py
+++ b/packages/opal-common/opal_common/sources/api_policy_source.py
@@ -46,6 +46,9 @@ class ApiPolicySource(BasePolicySource):
         token (str, optional):  auth token to include in connections to bundle server. Defaults to POLICY_BUNDLE_SERVER_TOKEN.
         token_id (str, optional):  auth token ID to include in connections to bundle server. Defaults to POLICY_BUNDLE_SERVER_TOKEN_ID.
         bundle_server_type (PolicyBundleServerType, optional):  the type of bundle server
+        region (str, optional): the aws region of s3 bucket containing the bundle
+        aws_role_arn (str, optional): the aws iam role to assume when accessing the s3 bucket. Only required when using temporary sts credentials.
+        aws_web_id_token_file (str, optional): the file containing a web id token for the target aws iam role. Only required when using temporary sts credentials.
     """
 
     def __init__(
@@ -56,8 +59,8 @@ def __init__(
         token: Optional[str] = None,
         token_id: Optional[str] = None,
         region: Optional[str] = None,
-        role_arn: Optional[str] = None,
-        token_file: Optional[str] = None,
+        aws_role_arn: Optional[str] = None,
+        aws_web_id_token_file: Optional[str] = None,
         bundle_server_type: Optional[PolicyBundleServerType] = None,
         policy_bundle_path=".",
         policy_bundle_git_add_pattern="*",
@@ -71,8 +74,8 @@ def __init__(
         self.token_id = token_id
         self.server_type = bundle_server_type
         self.region = region
-        self.role_arn = role_arn
-        self.token_file = token_file
+        self.aws_role_arn = aws_role_arn
+        self.aws_web_id_token_file = aws_web_id_token_file
         self.bundle_hash = None
         self.etag = None
         self.tmp_bundle_path = Path(policy_bundle_path)
@@ -135,11 +138,23 @@ async def api_update_policy(self) -> Tuple[bool, str, str]:
 
     @async_time_cache(ttl=3000)
     async def get_temporary_sts_credentials(self) -> tuple[str, str, str]:
-        assert self.token_file
-        assert self.role_arn
+        """
+        This function will fetch a set of temporary credentials for a IAM role
+        from Amazon STS. It requires an aws region, the arn for the target role
+        and the file containing the web token.
+
+        This function will return the id and secret key required for login.
+        When using temporary credentials, AWS also requires a session token
+        which this function also provides.
+
+        This result of this funciton is cached to avoid being rate limited by
+        STS.
+        """
+        assert self.aws_web_id_token_file
+        assert self.aws_role_arn
         assert self.region
 
-        async with aiofiles.open(self.token_file) as token_file:
+        async with aiofiles.open(self.aws_web_id_token_file) as token_file:
             token = await token_file.read()
 
         sts_url = f"sts.{self.region}.amazonaws.com"
@@ -147,7 +162,7 @@ async def get_temporary_sts_credentials(self) -> tuple[str, str, str]:
             "Action": "AssumeRoleWithWebIdentity",
             "DurationSeconds": "3600",
             "RoleSessionName": "Opal",
-            "RoleArn": self.role_arn,
+            "RoleArn": self.aws_role_arn,
             "WebIdentityToken": token,
             "Version": "2011-06-15",
         }
@@ -208,7 +223,7 @@ async def build_auth_headers(self, token=None, path=None):
             and token is not None
             and self.token_id is not None
         ):
-            logger.info("Using provided token to login to AWS_S3")
+            logger.info("Using provided token to log in to AWS_S3")
 
             split_url = urlparse(self.remote_source_url)
             host = split_url.netloc
@@ -219,11 +234,11 @@ async def build_auth_headers(self, token=None, path=None):
             )
         elif (
             self.server_type == PolicyBundleServerType.AWS_S3
-            and self.role_arn is not None
-            and self.token_file is not None
+            and self.aws_role_arn is not None
+            and self.aws_web_id_token_file is not None
             and self.region is not None
         ):
-            logger.info("Using IAM Web auth to login to AWS_S3")
+            logger.info("Using IAM Web auth to log in to AWS_S3")
 
             split_url = urlparse(self.remote_source_url)
             host = split_url.netloc
@@ -370,4 +385,4 @@ async def check_for_changes(self):
                 prev_head=prev,
                 new_head=latest,
             )
-            await self._on_new_policy(old=prev_commit, new=new_commit)
+            await self._on_new_policy(old=prev_commit, new=new_commit)
\ No newline at end of file
diff --git a/packages/opal-common/opal_common/utils.py b/packages/opal-common/opal_common/utils.py
index 4f9e1f56e..f0e86fbed 100644
--- a/packages/opal-common/opal_common/utils.py
+++ b/packages/opal-common/opal_common/utils.py
@@ -296,7 +296,7 @@ def run_coro(self, coro: Coroutine):
 
 def async_time_cache(ttl: float):
     """
-    This decorator a wrapper around lru_cache that makes it time sensitive.
+    This decorator is a wrapper around lru_cache that makes it time sensitive.
 
     ttl is in seconds
     """
@@ -316,4 +316,4 @@ def ret(*args, **kwargs):
 
         return ret
 
-    return decorator
\ No newline at end of file
+    return decorator
diff --git a/packages/opal-server/opal_server/config.py b/packages/opal-server/opal_server/config.py
index 62cb91ada..d692fa90f 100644
--- a/packages/opal-server/opal_server/config.py
+++ b/packages/opal-server/opal_server/config.py
@@ -8,7 +8,6 @@
 from opal_common.schemas.webhook import GitWebhookRequestParams
 
 confi = Confi(prefix="OPAL_")
-# confi_no_prefix = Confi(prefix="")
 
 
 class PolicySourceTypes(str, Enum):
@@ -135,16 +134,18 @@ class OpalServerConfig(Confi):
         "us-east-1",
         description="The AWS region of the S3 bucket",
     )
-    # POLICY_BUNDLE_AWS_ROLE_ARN = confi_no_prefix.str(
-    #     "AWS_ROLE_ARN",
-    #     None,
-    #     description="The IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
-    # )
-    # POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE = confi_no_prefix.str(
-    #     "AWS_WEB_IDENTITY_TOKEN_FILE",
-    #     None,
-    #     description="The oidc token for the IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS",
-    # )
+    POLICY_BUNDLE_AWS_ROLE_ARN = confi.str(
+        "AWS_ROLE_ARN",
+        # default to the env var injected by aws
+        os.getenv("AWS_ROLE_ARN"),
+        description="The IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS, but can be overridden if required.",
+    )
+    POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE = confi.str(
+        "AWS_WEB_IDENTITY_TOKEN_FILE",
+        # default to the env var injected by aws
+        os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE"),
+        description="The oidc token for the IAM role to be used when accessing the bundle server. This is set by AWS automatically in EKS, but can be overridden if required.",
+    )
     POLICY_BUNDLE_TMP_PATH = confi.str(
         "POLICY_BUNDLE_TMP_PATH",
         "/tmp/bundle.tar.gz",
@@ -329,4 +330,4 @@ def on_load(self):
             self.SERVER_BIND_PORT = int(self.SERVER_PORT)
 
 
-opal_server_config = OpalServerConfig(prefix="OPAL_")
\ No newline at end of file
+opal_server_config = OpalServerConfig(prefix="OPAL_")
diff --git a/packages/opal-server/opal_server/policy/watcher/factory.py b/packages/opal-server/opal_server/policy/watcher/factory.py
index 7b3af1ae4..e08d7a997 100644
--- a/packages/opal-server/opal_server/policy/watcher/factory.py
+++ b/packages/opal-server/opal_server/policy/watcher/factory.py
@@ -130,10 +130,8 @@ def setup_watcher_task(
             policy_bundle_path=opal_server_config.POLICY_BUNDLE_TMP_PATH,
             policy_bundle_git_add_pattern=opal_server_config.POLICY_BUNDLE_GIT_ADD_PATTERN,
             region=policy_bundle_aws_region,
-            # role_arn=opal_server_config.POLICY_BUNDLE_AWS_ROLE_ARN,
-            # token_file=opal_server_config.POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE,
-            role_arn=os.getenv("AWS_ROLE_ARN"),
-            token_file=os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE"),
+            aws_role_arn=opal_server_config.POLICY_BUNDLE_AWS_ROLE_ARN,
+            aws_web_id_token_file=opal_server_config.POLICY_BUNDLE_AWS_WEB_IDENTITY_TOKEN_FILE,
         )
     else:
         raise ValueError("Unknown value for OPAL_POLICY_SOURCE_TYPE")
@@ -145,4 +143,4 @@ def setup_watcher_task(
             bundle_ignore=bundle_ignore,
         )
     )
-    return PolicyWatcherTask(watcher, pubsub_endpoint)
\ No newline at end of file
+    return PolicyWatcherTask(watcher, pubsub_endpoint)