singlestore-auth-iam 0.4.0__py3-none-any.whl

s2iam/__init__.py

@@ -0,0 +1,35 @@
+"""
+SingleStore Auth IAM - Python Client Library
+
+A Python client library for cloud provider identity detection and authentication
+with SingleStore's IAM service.
+"""
+
+__version__ = "0.4.0"
+
+from .api import DETECT_PROVIDER_DEFAULT_TIMEOUT, detect_provider
+from .jwt import get_jwt, get_jwt_api, get_jwt_database
+from .models import (
+    AssumeRoleNotSupported,
+    CloudIdentity,
+    CloudProviderNotFound,
+    CloudProviderType,
+    JWTType,
+    ProviderIdentityUnavailable,
+    ProviderNotDetected,
+)
+
+__all__ = [
+    "detect_provider",
+    "DETECT_PROVIDER_DEFAULT_TIMEOUT",
+    "get_jwt",
+    "get_jwt_database",
+    "get_jwt_api",
+    "CloudIdentity",
+    "CloudProviderType",
+    "JWTType",
+    "CloudProviderNotFound",
+    "ProviderNotDetected",
+    "ProviderIdentityUnavailable",
+    "AssumeRoleNotSupported",
+]

s2iam/api.py

@@ -0,0 +1,300 @@
+"""
+Main API for the s2iam library.
+"""
+
+import asyncio
+import os
+import queue
+import threading
+import time
+from typing import Any, Dict, List, NoReturn, Optional
+
+from .aws import new_client as new_aws_client
+from .azure import new_client as new_azure_client
+from .gcp import new_client as new_gcp_client
+from .models import CloudProviderClient, CloudProviderNotFound, Logger
+
+DETECT_PROVIDER_DEFAULT_TIMEOUT: float = 10.0
+"""Default timeout (seconds) for provider detection.
+
+Rationale: Prefer avoiding false negatives over minimizing worst‑case wait.
+10s ceiling handles slow or throttled metadata on constrained CI VMs while
+early success still returns in sub‑second typical cases. Mirrors Go/Java parity.
+"""
+
+
+class DefaultLogger:
+    """Default logger implementation."""
+
+    def log(self, message: str) -> None:
+        """Log a message to stdout."""
+        print(f"[s2iam] {message}")
+
+
+async def detect_provider(
+    timeout: float = DETECT_PROVIDER_DEFAULT_TIMEOUT,
+    logger: Optional[Logger] = None,
+    clients: Optional[list[CloudProviderClient]] = None,
+) -> CloudProviderClient:
+    """
+    Detect which cloud provider we're running on.
+
+    Args:
+        timeout: Detection timeout in seconds
+        logger: Optional logger instance
+        clients: Optional list of custom provider clients
+
+    Returns:
+        CloudProviderClient for the detected provider
+
+    Raises:
+        CloudProviderNotFound: If no provider can be detected
+    """
+    # Set up logger only if explicit debugging flag is set; production code must not branch
+    # on test harness-only environment variables. Rich diagnostics are instead
+    # surfaced via aggregated exception messages below.
+    debugging = os.environ.get("S2IAM_DEBUGGING", "").lower() == "true"
+    debug_timing = os.environ.get("S2IAM_DEBUG_TIMING", "").lower() == "true"
+    if logger is None and debugging:
+        logger = DefaultLogger()
+
+    # Create default clients if none provided
+    if clients is None:
+        clients = [
+            new_aws_client(logger),
+            new_gcp_client(logger),
+            new_azure_client(logger),
+        ]
+
+    # Phase 1: fast_detect sequentially (purely local; must not do network I/O).
+    for c in clients:
+        try:
+            await c.fast_detect()
+            if logger:
+                logger.log(f"Fast detected provider: {c.get_type().value}")
+            return c
+        except Exception:
+            # Not detected via fast path; move to next provider.
+            continue
+
+    # Phase 2: full detection using threads (mirrors Go goroutines + first-winner channel).
+    # Invariants relied upon here: each client's detect() MUST raise on negative outcome; only a
+    # positively detected client (internal flag set) returns normally. This prevents selecting a
+    # provider that will later fail with ProviderNotDetected when building identity headers.
+    result_queue: "queue.Queue[CloudProviderClient]" = queue.Queue()
+    stop_event = threading.Event()
+    all_errors: list[str] = []
+    errors_lock = threading.Lock()
+    # Structured per-provider status for enhanced error reporting.
+    # Each element: {provider, status=success|error|timeout|skipped, elapsed_ms?, error?}
+    provider_status: List[Dict[str, Any]] = []
+    status_lock = threading.Lock()
+
+    def record_status(entry: Dict[str, Any]) -> None:
+        with status_lock:
+            provider_status.append(entry)
+
+    # Track per-thread event loops so we can cancel/stop them on global timeout
+    provider_loops: list[asyncio.AbstractEventLoop] = []
+    loops_lock = threading.Lock()
+
+    def test_provider_sync(client: CloudProviderClient) -> None:
+        """Test a provider in a thread (like Go goroutine)."""
+        if stop_event.is_set():
+            record_status({"provider": client.get_type().value, "status": "skipped"})
+            return
+        thread_start = time.monotonic()
+        if logger and (debugging or debug_timing):
+            logger.log(f"DETECT_THREAD_START provider={client.get_type().value} outer_timeout_s={timeout}")
+        try:
+            # Run the async detect() in this thread's event loop
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            with loops_lock:
+                provider_loops.append(loop)
+            try:
+                detect_coro = client.detect()
+                loop.run_until_complete(detect_coro)
+                # Post-call validation: provider must internally mark detected.
+                is_detected = True
+                try:
+                    # Allow provider to expose explicit method; fallback to attribute.
+                    if hasattr(client, "is_detected") and callable(getattr(client, "is_detected")):
+                        is_detected = bool(client.is_detected())
+                    elif hasattr(client, "_detected"):
+                        is_detected = bool(getattr(client, "_detected"))
+                except Exception:  # noqa: BLE001 - conservative: treat as success only if flag readable
+                    pass
+                elapsed_ms = int((time.monotonic() - thread_start) * 1000)
+                if is_detected and not stop_event.is_set():
+                    result_queue.put(client)
+                    stop_event.set()
+                    record_status(
+                        {
+                            "provider": client.get_type().value,
+                            "status": "success",
+                            "elapsed_ms": elapsed_ms,
+                        }
+                    )
+                    if logger and (debugging or debug_timing):
+                        logger.log(f"DETECT_THREAD_SUCCESS provider={client.get_type().value} elapsed_ms={elapsed_ms}")
+                else:
+                    record_status(
+                        {
+                            "provider": client.get_type().value,
+                            "status": "not_detected",
+                            "elapsed_ms": elapsed_ms,
+                        }
+                    )
+                    if logger and (debugging or debug_timing):
+                        logger.log(
+                            f"DETECT_THREAD_NOT_DETECTED provider={client.get_type().value} elapsed_ms={elapsed_ms}"
+                        )
+            finally:
+                loop.close()
+        except Exception as e:
+            with errors_lock:
+                all_errors.append(f"Provider {client.get_type().value} detection failed: {e}")
+            elapsed_ms = int((time.monotonic() - thread_start) * 1000)
+            record_status(
+                {
+                    "provider": client.get_type().value,
+                    "status": "error",
+                    "elapsed_ms": elapsed_ms,
+                    "error": str(e)[:400],
+                }
+            )
+            if logger and (debugging or debug_timing):
+                logger.log(f"DETECT_THREAD_ERROR provider={client.get_type().value} elapsed_ms={elapsed_ms} error={e}")
+
+    # Start threads for each provider (like Go goroutines)
+    threads = []
+    for client in clients:
+        thread = threading.Thread(target=test_provider_sync, args=(client,))
+        thread.daemon = True
+        thread.start()
+        threads.append(thread)
+
+    # Wait for first result or timeout (like Go select)
+    detection_start = time.monotonic()
+    # Track how many threads have finished (success or error) to allow early exit when all done.
+    total_clients = len(clients)
+    # Poll loop instead of single blocking get so we can detect early-failure condition.
+    remaining = timeout
+    interval = 0.05  # 50ms poll granularity (balance: coarse enough to reduce wakeups, fine enough for fast success)
+    while remaining > 0:
+        start_poll = time.monotonic()
+        try:
+            result: CloudProviderClient = result_queue.get(timeout=min(interval, remaining))
+            stop_event.set()  # Ensure all threads stop
+            total_elapsed_ms = int((time.monotonic() - detection_start) * 1000)
+            if logger:
+                if debugging or debug_timing:
+                    logger.log(
+                        (
+                            "DETECT_COMPLETE status=success "
+                            f"provider={result.get_type().value} "
+                            f"total_elapsed_ms={total_elapsed_ms}"
+                        )
+                    )
+                else:
+                    logger.log(f"Detected provider: {result.get_type().value}")
+            return result
+        except queue.Empty:
+            pass
+        # Early failure: if every thread has produced a terminal status (success/error/timeout/skipped)
+        with status_lock:
+            finished = sum(1 for ps in provider_status if ps["status"] in {"success", "error", "timeout", "skipped"})
+        if finished >= total_clients and not any(ps["status"] == "success" for ps in provider_status):
+            # All threads ended without success -> raise immediately (no need to wait remaining timeout)
+            break
+        remaining -= time.monotonic() - start_poll
+    else:
+        # Loop ended naturally (remaining <= 0) without a success result
+        if logger and (debugging or debug_timing):
+            logger.log(
+                "DETECT_LOOP_COMPLETE no-success reason=timeout-before-result "
+                f"elapsed_ms={int((time.monotonic()-detection_start)*1000)}"
+            )
+
+    # No provider detected within timeout or all failed fast.
+    stop_event.set()  # Ensure all threads stop
+
+    # If we broke out of loop without returning (early failure or timeout), raise composed error.
+    _raise_detection_timeout(
+        timeout=timeout,
+        detection_start=detection_start,
+        clients=clients,
+        threads=threads,
+        provider_status=provider_status,
+        status_lock=status_lock,
+        all_errors=all_errors,
+        logger=logger,
+        debugging=debugging,
+        debug_timing=debug_timing,
+        stop_event=stop_event,
+        provider_loops=provider_loops,
+        loops_lock=loops_lock,
+    )
+
+
+def _raise_detection_timeout(
+    *,
+    timeout: float,
+    detection_start: float,
+    clients: list[CloudProviderClient],
+    threads: list[threading.Thread],
+    provider_status: list[dict[str, Any]],
+    status_lock: threading.Lock,
+    all_errors: list[str],
+    logger: Optional[Logger],
+    debugging: bool,
+    debug_timing: bool,
+    stop_event: threading.Event,
+    provider_loops: list[asyncio.AbstractEventLoop],
+    loops_lock: threading.Lock,
+) -> NoReturn:
+    """Compose and raise CloudProviderNotFound for a detection timeout.
+
+    Isolated to keep the main detect_provider flow skimmable and ease future
+    experiments (e.g., per-provider granular timeouts or retry policy integration).
+    """
+    stop_event.set()
+    # Attempt to stop any active provider event loops to prevent post-timeout drift
+    with loops_lock:
+        for loop in provider_loops:
+            if loop.is_running():
+                try:
+                    loop.call_soon_threadsafe(loop.stop)
+                except Exception:  # noqa: BLE001 - best effort cancellation
+                    pass
+    for thread in threads:
+        # 50ms join chosen: keeps worst-case added delay bounded (< provider timeout granularity)
+        # while giving threads a chance to observe stop_event and exit cleanly.
+        thread.join(timeout=0.05)
+    total_elapsed_ms = int((time.monotonic() - detection_start) * 1000)
+    with status_lock:
+        known = {p["provider"] for p in provider_status}
+        for c in clients:
+            name = c.get_type().value
+            if name not in known:
+                provider_status.append({"provider": name, "status": "timeout"})
+    if logger and (debugging or debug_timing):
+        joined_errors_dbg = " | ".join(all_errors)[:800]
+        logger.log(
+            (
+                "DETECT_COMPLETE status=timeout "
+                f"total_elapsed_ms={total_elapsed_ms} timeout_s={timeout} "
+                f"errors='{joined_errors_dbg}'"
+            )
+        )
+    summary = ", ".join(
+        f"{ps['provider']}:{ps['status']}{('@'+str(ps['elapsed_ms'])+'ms') if 'elapsed_ms' in ps else ''}"
+        for ps in provider_status
+    )
+    errors_str = " | ".join(all_errors)[:800] if all_errors else "<no-provider-errors>"
+    raise CloudProviderNotFound(
+        "Provider detection timed out: "
+        f"timeout_s={timeout} total_elapsed_ms={total_elapsed_ms} providers={len(clients)} "
+        f"error_count={len(all_errors)} provider_status=[{summary}] errors=[{errors_str}]"
+    )

s2iam/aws/__init__.py

@@ -0,0 +1,271 @@
+"""AWS cloud provider client implementation.
+
+Single implementation aligned with the Go reference: fast env/IMDS detection,
+STS fallback, optional role assumption, region derivation, and identity header
+generation.
+"""
+
+import asyncio
+import os
+from typing import Any, Optional
+
+from ..models import (
+    CloudIdentity,
+    CloudProviderClient,
+    CloudProviderType,
+    Logger,
+    ProviderIdentityUnavailable,
+    ProviderNotDetected,
+)
+
+
+class AWSClient(CloudProviderClient):
+    _logger: Optional[Logger]
+    _detected: bool
+    _region: Optional[str]
+    _identity: Optional[CloudIdentity]
+    _role_arn: Optional[str]
+    _sts_client: Optional[Any]
+    _session: Optional[Any]
+
+    def __init__(self, logger: Optional[Logger] = None):
+        self._logger = logger
+        self._detected = False
+        self._region = None
+        self._identity = None
+        self._role_arn = None
+        self._sts_client = None
+        self._session = None
+
+    def _log(self, message: str) -> None:
+        if self._logger:
+            self._logger.log(f"AWS: {message}")
+
+    async def _check_metadata_service(self) -> bool:
+        """Best-effort IMDSv2 then IMDSv1 probe (<= ~3s worst case)."""
+        try:  # noqa: BLE001
+            import aiohttp
+
+            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=3)) as session:
+                async with session.put(
+                    "http://169.254.169.254/latest/api/token",
+                    headers={"X-aws-ec2-metadata-token-ttl-seconds": "21600"},
+                ) as token_resp:
+                    if token_resp.status == 200:
+                        token = await token_resp.text()
+                        async with session.get(
+                            "http://169.254.169.254/latest/meta-data/instance-id",
+                            headers={"X-aws-ec2-metadata-token": token},
+                        ) as resp:
+                            if resp.status == 200:
+                                return True
+
+                async with session.get(
+                    "http://169.254.169.254/latest/meta-data/instance-id",
+                    timeout=aiohttp.ClientTimeout(total=2),
+                ) as resp:
+                    return resp.status == 200
+        except Exception as e:  # noqa: BLE001
+            self._log(f"Metadata service check failed: {e}")
+            return False
+
+    async def detect(self) -> None:
+        # Full (network-inclusive) detection. Raise on failure so orchestrator never
+        # selects an undetected client (prevents later ProviderNotDetected errors).
+        self._log("Starting AWS detection (full phase)")
+        if self._detected:
+            return
+
+        if await self._check_metadata_service():
+            self._detected = True
+            self._log("Detected via metadata service")
+            return
+
+        try:
+            import boto3  # optional dependency in some usage contexts
+        except ImportError as e:  # noqa: BLE001
+            self._log(f"boto3 import failed: {e}")
+        else:
+            try:  # noqa: BLE001
+                sts_client = boto3.client("sts")
+                identity = sts_client.get_caller_identity()
+                if identity.get("Account"):
+                    self._detected = True
+                    self._log("Detected via STS")
+                    return
+            except Exception as e:  # noqa: BLE001
+                self._log(f"STS detection failed: {e}")
+
+        self._log("AWS full detection did not succeed; raising")
+        raise Exception("AWS provider not detected")
+
+    async def fast_detect(self) -> None:
+        """Fast detection: env only, no network calls."""
+        # IRSA / web identity short-circuit: ONLY honor explicit env vars.
+        if os.environ.get("AWS_WEB_IDENTITY_TOKEN_FILE") or os.environ.get("AWS_ROLE_ARN"):
+            self._detected = True
+            self._log("FastDetect: IRSA environment variables present")
+            return
+
+        for var in ("AWS_EXECUTION_ENV", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_LAMBDA_FUNCTION_NAME"):
+            if os.environ.get(var):
+                self._detected = True
+                self._log(f"FastDetect: detected via env var {var}")
+                return
+        raise Exception("FastDetect: no AWS indicators")
+
+    async def _ensure_region(self) -> None:
+        if self._region:
+            return
+
+        region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
+        if not region:
+            try:  # noqa: BLE001
+                import aiohttp
+
+                async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=3)) as session:
+                    async with session.put(
+                        "http://169.254.169.254/latest/api/token",
+                        headers={"X-aws-ec2-metadata-token-ttl-seconds": "21600"},
+                    ) as token_resp:
+                        if token_resp.status == 200:
+                            token = await token_resp.text()
+                            async with session.get(
+                                "http://169.254.169.254/latest/meta-data/placement/region",
+                                headers={"X-aws-ec2-metadata-token": token},
+                            ) as region_resp:
+                                if region_resp.status == 200:
+                                    region = await region_resp.text()
+                                    self._log(f"Region from metadata: {region}")
+            except Exception as e:  # noqa: BLE001
+                self._log(f"Region metadata lookup failed: {e}")
+
+        if not region:
+            region = "us-east-1"
+            self._log("Defaulting region to us-east-1")
+        self._region = region
+
+    def get_type(self) -> CloudProviderType:
+        return CloudProviderType.AWS
+
+    def assume_role(self, role_identifier: str) -> "AWSClient":
+        clone = AWSClient(self._logger)
+        clone._detected = self._detected
+        clone._region = self._region
+        clone._role_arn = role_identifier
+        clone._sts_client = self._sts_client
+        return clone
+
+    async def get_identity_headers(
+        self, additional_params: Optional[dict[str, str]] = None
+    ) -> tuple[dict[str, str], CloudIdentity]:  # noqa: D401,E501
+        if not self._detected:
+            raise ProviderNotDetected("AWS provider not detected, call detect() first")
+
+        if not self._sts_client:
+            import boto3
+
+            await self._ensure_region()
+            self._session = boto3.Session()
+            self._sts_client = self._session.client("sts", region_name=self._region)
+            self._log("Initialized STS client")
+
+        if self._sts_client is None or self._session is None:
+            # Can happen when clone created via assume_role (sts_client copied but session not)
+            import boto3
+
+            if self._session is None:
+                self._session = boto3.Session()
+            if self._sts_client is None:
+                self._sts_client = self._session.client("sts", region_name=self._region)
+            self._log("Recovered missing STS session/client state")
+
+        # Narrow optionals after explicit check
+        sts_client = self._sts_client
+        session_obj = self._session
+
+        loop = asyncio.get_event_loop()
+
+        try:  # noqa: BLE001
+            if self._role_arn:
+                self._log(f"Assuming role {self._role_arn}")
+                assume_resp = await loop.run_in_executor(
+                    None,
+                    lambda: sts_client.assume_role(
+                        RoleArn=self._role_arn,
+                        RoleSessionName="s2iam-session",
+                    ),
+                )
+                creds = assume_resp["Credentials"]
+                import boto3
+
+                assumed_session = boto3.Session(
+                    aws_access_key_id=creds["AccessKeyId"],
+                    aws_secret_access_key=creds["SecretAccessKey"],
+                    aws_session_token=creds["SessionToken"],
+                    region_name=self._region,
+                )
+                assumed_sts = assumed_session.client("sts")
+                identity_resp = await loop.run_in_executor(None, assumed_sts.get_caller_identity)
+                headers = {
+                    "X-AWS-Access-Key-ID": creds["AccessKeyId"],
+                    "X-AWS-Secret-Access-Key": creds["SecretAccessKey"],
+                    "X-AWS-Session-Token": creds["SessionToken"],
+                }
+            else:
+                identity_resp = await loop.run_in_executor(None, sts_client.get_caller_identity)
+                role_assumed = (
+                    ":assumed-role/" in identity_resp["Arn"] or os.environ.get("AWS_SESSION_TOKEN") is not None
+                )
+                if role_assumed:
+                    creds = session_obj.get_credentials()
+                    headers = {
+                        "X-AWS-Access-Key-ID": creds.access_key,
+                        "X-AWS-Secret-Access-Key": creds.secret_key,
+                        "X-Cloud-Provider": "aws",
+                    }
+                    if creds.token:
+                        headers["X-AWS-Session-Token"] = creds.token
+                    if os.environ.get("AWS_WEB_IDENTITY_TOKEN_FILE") or os.environ.get("AWS_ROLE_ARN"):
+                        self._log("Using IRSA web identity session credentials")
+                else:
+                    self._log("Getting session token for static credentials")
+                    session_resp = await loop.run_in_executor(None, sts_client.get_session_token)
+                    sc = session_resp["Credentials"]
+                    headers = {
+                        "X-AWS-Access-Key-ID": sc["AccessKeyId"],
+                        "X-AWS-Secret-Access-Key": sc["SecretAccessKey"],
+                        "X-AWS-Session-Token": sc["SessionToken"],
+                        "X-Cloud-Provider": "aws",
+                    }
+
+            arn = identity_resp["Arn"]
+            parts = arn.split(":")
+            region_from_arn = parts[3] if len(parts) > 3 else ""
+            resource_type = ""
+            if len(parts) > 5:
+                res_parts = parts[5].split("/")
+                if res_parts and res_parts[0]:
+                    resource_type = res_parts[0]
+
+            # If region unset locally (IRSA path without env/metadata), adopt ARN region
+            if not self._region and region_from_arn:
+                self._region = region_from_arn
+                self._log(f"Derived region from ARN: {self._region}")
+
+            identity = CloudIdentity(
+                provider=CloudProviderType.AWS,
+                identifier=arn,
+                account_id=identity_resp["Account"],
+                region=region_from_arn,
+                resource_type=resource_type,
+            )
+            self._log(f"Generated headers for identity: {identity.identifier}")
+            return headers, identity
+        except Exception as e:  # noqa: BLE001
+            self._log(f"Failed to build identity headers: {e}")
+            raise ProviderIdentityUnavailable(f"Failed to get AWS identity: {e}")
+
+
+def new_client(logger: Optional[Logger] = None) -> CloudProviderClient:
+    return AWSClient(logger)