[ROB-2024] Irsa cross account (#33)

* add-cross-account-support

* made it a bit cleaner

* safer code

* bump version

* added arn to prometheus config

* add assume role to custom connect

* updates, compatibility with newer pydantics

Author: Avi-Robusta

poetry.lock

@@ -1,15 +1,20 @@
+import os
 from datetime import datetime
-from typing import Any, Dict, Optional, List
+from typing import Any, Dict, Optional
+import logging
 
 import requests
 import boto3
 from botocore.auth import S3SigV4Auth
 from botocore.awsrequest import AWSRequest
 from botocore.credentials import Credentials
 from prometheus_api_client import PrometheusApiClientException
+from botocore.exceptions import BotoCoreError, ClientError
 
 from prometrix.connect.custom_connect import CustomPrometheusConnect
 
+SA_TOKEN_PATH = os.environ.get("SA_TOKEN_PATH", "/var/run/secrets/eks.amazonaws.com/serviceaccount/token")
+AWS_ASSUME_ROLE = os.environ.get("AWS_ASSUME_ROLE")
 
 class AWSPrometheusConnect(CustomPrometheusConnect):
     def __init__(
@@ -19,6 +24,7 @@ def __init__(
         region: str,
         service_name: str,
         token: Optional[str] = None,
+        assume_role_arn: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -36,6 +42,45 @@ def __init__(
                 raise RuntimeError("No AWS credentials found (neither static keys nor IRSA)")
             self._credentials = creds
 
+        role_to_assume = assume_role_arn or AWS_ASSUME_ROLE
+        if role_to_assume:
+            self._assume_role_with_web_identity(role_to_assume)
+
+    def _assume_role_with_web_identity(self, role_arn: str) -> None:
+        """Assume the given role using the pod's service-account web identity token."""
+        try:
+            with open(SA_TOKEN_PATH, "r", encoding="utf-8") as f:
+                web_identity_token = f.read().strip()
+        except FileNotFoundError:
+            raise RuntimeError(f"Service Account token not found at {SA_TOKEN_PATH}")
+
+        try:
+            sts = boto3.client("sts", region_name=self.region)
+            resp = sts.assume_role_with_web_identity(
+                RoleArn=role_arn,
+                RoleSessionName="amp-auto",
+                WebIdentityToken=web_identity_token,
+            )
+
+            credentials = resp.get("Credentials")
+            if not credentials:
+                logging.error("Invalid assume role response {resp}")
+                return
+
+            required_fields = ["AccessKeyId", "SecretAccessKey", "SessionToken"]
+            missing = [f for f in required_fields if not credentials.get(f)]
+            if missing:
+                logging.error("Missing required credential fields: {missing}. Raw response: {resp}")
+                raise Exception(f"Failed to assume role: missing fields {missing}")
+
+            self._credentials = Credentials(
+                credentials["AccessKeyId"],credentials["SecretAccessKey"], credentials["SessionToken"]
+            )
+        except (ClientError, BotoCoreError, Exception) as e:
+            raise Exception(
+                f"Failed to assume role {role_arn} with web identity: {str(e)}"
+            )
+
     def _build_auth(self) -> S3SigV4Auth:
         """Builds fresh SigV4 auth with current credentials (handles rotation)."""
         frozen = self._credentials.get_frozen_credentials()
@@ -53,6 +98,7 @@ def signed_request(
             headers=dict(request.headers),
             verify=verify,
             data=data,
+            params=params,
         )
 
     def _custom_query(self, query: str, params: dict = None):

prometrix/connect/aws_connect.py

@@ -1,7 +1,12 @@
 from enum import Enum
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel, SecretStr
+try:
+    # Works if Pydantic v2 is installed
+    from pydantic.v1 import BaseModel, SecretStr
+except ImportError:
+    # Fallback if running under Pydantic v1
+    from pydantic import BaseModel, SecretStr
 
 
 class PrometheusApis(Enum):
@@ -35,6 +40,7 @@ class AWSPrometheusConfig(PrometheusConfig):
     token: Optional[str] = None
     service_name: str = "aps"
     aws_region: str
+    assume_role_arn: Optional[str] = None
     supported_apis: List[PrometheusApis] = [
         PrometheusApis.QUERY,
         PrometheusApis.QUERY_RANGE,

prometrix/models/prometheus_config.py

@@ -36,6 +36,7 @@ def get_custom_prometheus_connect(
             secret_key=prom_config.secret_access_key,
             service_name=prom_config.service_name,
             region=prom_config.aws_region,
+            assume_role_arn=prom_config.assume_role_arn,
             config=prom_config,
             token=prom_config.token,
         )

prometrix/utils.py

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "prometrix"
-version = "0.2.2"
+version = "0.2.3"
 authors = ["Avi Kotlicky <[email protected]>"]
 readme = "README.md"
 packages = [{include = "prometrix"}]
@@ -14,7 +14,7 @@ description = "A Python Prometheus client for all Prometheus instances."
 python = "^3.8"
 boto3 = "^1.28.15"
 botocore = "^1.31.15"
-pydantic = "^1.8.1"
+pydantic = ">=1.8.1,<3"
 prometheus-api-client = "^0.5.3"
 pillow = "^10.3.0" # added to Pin transitive dependency, not needed directly
 fonttools = "^4.43.0" # added to Pin transitive dependency, not needed directly

pyproject.toml

@@ -5,7 +5,12 @@
 
 import yaml
 from prometheus_api_client import PrometheusApiClientException
-from pydantic import ValidationError
+
+try:
+    from pydantic.v1 import ValidationError
+except ImportError:
+    from pydantic import ValidationError
+
 from pytimeparse.timeparse import timeparse
 
 from prometrix import (AWSPrometheusConfig, AzurePrometheusConfig,