nui-python-shared-utils 1.3.0__py3-none-any.whl

nui_lambda_shared_utils/es_query_builder.py

@@ -0,0 +1,315 @@
+"""
+Elasticsearch query builder utilities for consistent query patterns across Lambda services.
+Provides helper functions for building common ES queries used in application monitoring.
+"""
+
+from typing import Dict, List, Optional, Union
+from datetime import datetime, timedelta
+
+
+class ESQueryBuilder:
+    """
+    Builder for creating Elasticsearch queries with common patterns.
+
+    Example:
+        builder = ESQueryBuilder()
+        query = (builder
+            .with_time_range(start_time, end_time)
+            .with_term('environment', 'prod')
+            .with_service('order-service')
+            .add_aggregation('error_rate', 'avg', 'response.status')
+            .build()
+        )
+    """
+
+    def __init__(self):
+        self.query = {"bool": {"must": [], "must_not": [], "should": [], "filter": []}}
+        self.aggregations = {}
+        self.size = 0  # Default to aggregation-only queries
+        self.sort = []
+
+    def with_time_range(self, start: datetime, end: datetime, field: str = "@timestamp") -> "ESQueryBuilder":
+        """Add time range filter."""
+        self.query["bool"]["must"].append(
+            {
+                "range": {
+                    field: {
+                        "gte": start.isoformat() if hasattr(start, "isoformat") else start,
+                        "lte": end.isoformat() if hasattr(end, "isoformat") else end,
+                    }
+                }
+            }
+        )
+        return self
+
+    def with_term(self, field: str, value: Union[str, int, bool]) -> "ESQueryBuilder":
+        """Add exact term match."""
+        self.query["bool"]["must"].append({"term": {field: value}})
+        return self
+
+    def with_terms(self, field: str, values: List[Union[str, int]]) -> "ESQueryBuilder":
+        """Add terms match (OR condition for multiple values)."""
+        self.query["bool"]["must"].append({"terms": {field: values}})
+        return self
+
+    def with_service(self, service: str) -> "ESQueryBuilder":
+        """Add service name filter."""
+        return self.with_term("service_name", service)
+
+    def with_environment(self, env: str = "prod") -> "ESQueryBuilder":
+        """Add environment filter."""
+        return self.with_term("environment", env)
+
+    def with_error_filter(self, min_status: int = 400) -> "ESQueryBuilder":
+        """Add filter for error responses."""
+        self.query["bool"]["must"].append({"range": {"response.status": {"gte": min_status}}})
+        return self
+
+    def exclude_pattern(self, field: str, pattern: str) -> "ESQueryBuilder":
+        """Exclude documents matching a pattern."""
+        self.query["bool"]["must_not"].append({"wildcard": {field: pattern}})
+        return self
+
+    def with_prefix(self, field: str, prefix: str) -> "ESQueryBuilder":
+        """Add prefix match."""
+        self.query["bool"]["must"].append({"prefix": {field: prefix}})
+        return self
+
+    def add_aggregation(self, name: str, agg_type: str, field: str, **kwargs) -> "ESQueryBuilder":
+        """Add a simple aggregation."""
+        self.aggregations[name] = {agg_type: {"field": field, **kwargs}}
+        return self
+
+    def add_date_histogram(
+        self, name: str, field: str = "@timestamp", interval: str = "5m", **kwargs
+    ) -> "ESQueryBuilder":
+        """Add date histogram aggregation."""
+        self.aggregations[name] = {
+            "date_histogram": {"field": field, "fixed_interval": interval, "min_doc_count": 0, **kwargs}
+        }
+        return self
+
+    def add_terms_aggregation(self, name: str, field: str, size: int = 10, **kwargs) -> "ESQueryBuilder":
+        """Add terms aggregation for top values."""
+        self.aggregations[name] = {"terms": {"field": field, "size": size, **kwargs}}
+        return self
+
+    def add_percentiles(self, name: str, field: str, percents: List[float] = None) -> "ESQueryBuilder":
+        """Add percentiles aggregation."""
+        if percents is None:
+            percents = [50, 95, 99]
+
+        self.aggregations[name] = {"percentiles": {"field": field, "percents": percents}}
+        return self
+
+    def add_nested_aggregation(self, parent_name: str, child_aggs: Dict) -> "ESQueryBuilder":
+        """Add nested aggregations."""
+        if parent_name not in self.aggregations:
+            raise ValueError(f"Parent aggregation '{parent_name}' not found")
+
+        self.aggregations[parent_name]["aggs"] = child_aggs
+        return self
+
+    def with_size(self, size: int) -> "ESQueryBuilder":
+        """Set number of documents to return."""
+        self.size = size
+        return self
+
+    def add_sort(self, field: str, order: str = "desc") -> "ESQueryBuilder":
+        """Add sort criteria."""
+        self.sort.append({field: {"order": order}})
+        return self
+
+    def build(self) -> Dict:
+        """Build the final query."""
+        query = {"query": self.query, "size": self.size}
+
+        if self.aggregations:
+            query["aggs"] = self.aggregations
+
+        if self.sort:
+            query["sort"] = self.sort
+
+        return query
+
+
+# Pre-built query templates for common patterns
+def build_error_rate_query(service: str, start_time: datetime, end_time: datetime, interval: str = "5m") -> Dict:
+    """Build query for service error rate over time."""
+    builder = ESQueryBuilder()
+    return (
+        builder.with_time_range(start_time, end_time)
+        .with_service(service)
+        .with_environment("prod")
+        .add_date_histogram("error_timeline", interval=interval)
+        .add_nested_aggregation(
+            "error_timeline",
+            {
+                "total_requests": {"value_count": {"field": "request.id"}},
+                "error_requests": {
+                    "filter": {"range": {"response.status": {"gte": 400}}},
+                    "aggs": {"count": {"value_count": {"field": "response.status"}}},
+                },
+                "error_rate": {
+                    "bucket_script": {
+                        "buckets_path": {"errors": "error_requests>count", "total": "total_requests"},
+                        "script": "params.errors / params.total * 100",
+                    }
+                },
+            },
+        )
+        .build()
+    )
+
+
+def build_top_errors_query(service: str, start_time: datetime, end_time: datetime, top_n: int = 10) -> Dict:
+    """Build query for top error messages."""
+    builder = ESQueryBuilder()
+    return (
+        builder.with_time_range(start_time, end_time)
+        .with_service(service)
+        .with_environment("prod")
+        .with_error_filter()
+        .add_terms_aggregation("top_errors", "error.keyword", size=top_n)
+        .add_nested_aggregation(
+            "top_errors",
+            {
+                "status_codes": {"terms": {"field": "response.status", "size": 5}},
+                "sample_error": {"top_hits": {"size": 1, "_source": ["error", "request.path", "@timestamp"]}},
+            },
+        )
+        .build()
+    )
+
+
+def build_response_time_query(service: str, start_time: datetime, end_time: datetime) -> Dict:
+    """Build query for response time metrics."""
+    builder = ESQueryBuilder()
+    return (
+        builder.with_time_range(start_time, end_time)
+        .with_service(service)
+        .with_environment("prod")
+        .add_percentiles("response_times", "response.time", [50, 90, 95, 99])
+        .add_aggregation("avg_response_time", "avg", "response.time")
+        .add_aggregation("max_response_time", "max", "response.time")
+        .add_date_histogram("response_timeline", interval="5m")
+        .add_nested_aggregation(
+            "response_timeline",
+            {
+                "avg_time": {"avg": {"field": "response.time"}},
+                "p95_time": {"percentiles": {"field": "response.time", "percents": [95]}},
+            },
+        )
+        .build()
+    )
+
+
+def build_service_volume_query(services: List[str], start_time: datetime, end_time: datetime) -> Dict:
+    """Build query for request volume across multiple services."""
+    builder = ESQueryBuilder()
+    return (
+        builder.with_time_range(start_time, end_time)
+        .with_terms("service_name", services)
+        .with_environment("prod")
+        .add_terms_aggregation("service_breakdown", "service_name", size=20)
+        .add_nested_aggregation(
+            "service_breakdown",
+            {
+                "request_count": {"value_count": {"field": "request.id"}},
+                "error_count": {
+                    "filter": {"range": {"response.status": {"gte": 400}}},
+                    "aggs": {"count": {"value_count": {"field": "response.status"}}},
+                },
+                "avg_response_time": {"avg": {"field": "response.time"}},
+            },
+        )
+        .build()
+    )
+
+
+def build_user_activity_query(start_time: datetime, end_time: datetime, user_field: str = "user.id") -> Dict:
+    """Build query for user activity metrics."""
+    builder = ESQueryBuilder()
+    return (
+        builder.with_time_range(start_time, end_time)
+        .with_environment("prod")
+        .add_aggregation("unique_users", "cardinality", user_field)
+        .add_terms_aggregation("top_users", user_field, size=20)
+        .add_nested_aggregation(
+            "top_users",
+            {
+                "request_count": {"value_count": {"field": "request.id"}},
+                "services_used": {"cardinality": {"field": "service_name"}},
+                "last_activity": {"max": {"field": "@timestamp"}},
+            },
+        )
+        .add_date_histogram("user_timeline", interval="1h")
+        .add_nested_aggregation("user_timeline", {"active_users": {"cardinality": {"field": user_field}}})
+        .build()
+    )
+
+
+def build_pattern_detection_query(
+    pattern: str, field: str = "message", start_time: Optional[datetime] = None, hours_back: int = 24
+) -> Dict:
+    """Build query to detect specific patterns in logs."""
+    if not start_time:
+        start_time = datetime.utcnow() - timedelta(hours=hours_back)
+    end_time = datetime.utcnow()
+
+    return {
+        "query": {
+            "bool": {
+                "must": [
+                    {"wildcard": {field: f"*{pattern}*"}},
+                    {"range": {"@timestamp": {"gte": start_time.isoformat(), "lte": end_time.isoformat()}}},
+                ]
+            }
+        },
+        "size": 100,
+        "sort": [{"@timestamp": {"order": "desc"}}],
+        "aggs": {
+            "service_breakdown": {"terms": {"field": "service_name", "size": 10}},
+            "timeline": {"date_histogram": {"field": "@timestamp", "fixed_interval": "1h", "min_doc_count": 1}},
+        },
+    }
+
+
+def build_tender_participant_query(tender_id: str, start_time: datetime, end_time: datetime) -> Dict:
+    """Build query for tender participant analysis."""
+    # Expand time range for pre-tender activity
+    expanded_start = start_time - timedelta(hours=2)
+    expanded_end = end_time + timedelta(minutes=30)
+
+    return {
+        "query": {
+            "bool": {
+                "must": [
+                    {"term": {"tender.id": tender_id}},
+                    {"range": {"created": {"gte": expanded_start.isoformat(), "lte": expanded_end.isoformat()}}},
+                ]
+            }
+        },
+        "size": 0,
+        "aggs": {
+            "participants": {
+                "terms": {"field": "division.id", "size": 100},
+                "aggs": {
+                    "participant_details": {"top_hits": {"size": 1, "_source": ["division.name", "company.name"]}},
+                    "bid_count": {"value_count": {"field": "created"}},
+                    "first_bid": {"min": {"field": "created"}},
+                    "last_bid": {"max": {"field": "created"}},
+                    "bid_values": {"stats": {"field": "price"}},
+                },
+            },
+            "total_bids": {"value_count": {"field": "created"}},
+            "bid_timeline": {
+                "date_histogram": {
+                    "field": "created",
+                    "fixed_interval": "1m",
+                    "min_doc_count": 0,
+                    "extended_bounds": {"min": start_time.isoformat(), "max": end_time.isoformat()},
+                }
+            },
+        },
+    }

nui_lambda_shared_utils/jwt_auth.py

@@ -0,0 +1,277 @@
+"""
+JWT validation utilities for AWS Lambda functions behind API Gateway.
+
+Uses RS256 signature verification with public keys stored in AWS Secrets Manager.
+Requires the `rsa` package (pure Python, ~100KB) — no PyJWT or cryptography needed at runtime.
+
+Install: pip install nui-python-shared-utils[jwt]
+"""
+
+import os
+import json
+import re
+import base64
+import time
+import logging
+from typing import TYPE_CHECKING, AbstractSet, Any, Dict, Optional, Tuple
+from urllib.parse import unquote
+
+from .secrets_helper import get_secret
+
+if TYPE_CHECKING:
+    import rsa
+
+log = logging.getLogger(__name__)
+
+JWT_CLOCK_SKEW_SECONDS = 30
+"""Tolerance in seconds for clock differences between token issuer and validator."""
+
+_rsa = None
+
+
+def _require_rsa():
+    """Lazy-import the rsa package, raising a clear error if not installed."""
+    global _rsa
+    if _rsa is None:
+        try:
+            import rsa
+
+            _rsa = rsa
+        except ImportError:
+            raise ImportError("The 'rsa' package is required for JWT support. Install with: pip install nui-python-shared-utils[jwt]") from None
+    return _rsa
+
+
+class JWTValidationError(Exception):
+    """Base exception for JWT validation failures."""
+
+    pass
+
+
+class AuthenticationError(JWTValidationError):
+    """Authentication failed — missing/invalid token or header."""
+
+    pass
+
+
+def get_jwt_public_key(secret_name: Optional[str] = None, key_field: str = "TOKEN_PUBLIC_KEY"):
+    """
+    Fetch PEM public key from AWS Secrets Manager and return as rsa.PublicKey.
+
+    Relies on secrets_helper cache — repeated calls with the same secret_name
+    do not make additional Secrets Manager API calls.
+
+    Args:
+        secret_name: Secrets Manager secret name. Falls back to JWT_PUBLIC_KEY_SECRET env var.
+        key_field: JSON field containing the PEM-encoded public key.
+
+    Returns:
+        rsa.PublicKey ready for signature verification.
+
+    Raises:
+        JWTValidationError: If secret or key field is missing/invalid.
+    """
+    rsa = _require_rsa()
+
+    secret_name = secret_name or os.environ.get("JWT_PUBLIC_KEY_SECRET")
+    if not secret_name:
+        raise JWTValidationError("No JWT public key secret name provided (set JWT_PUBLIC_KEY_SECRET or pass secret_name)")
+
+    secret = get_secret(secret_name)
+
+    pem_str = secret.get(key_field)
+    if not pem_str:
+        raise JWTValidationError(f"Field '{key_field}' not found in secret '{secret_name}'")
+
+    try:
+        return rsa.PublicKey.load_pkcs1_openssl_pem(pem_str.encode("utf-8"))
+    except Exception as e:
+        raise JWTValidationError(
+            f"Failed to load public key from '{key_field}' — expected PKCS#8 PEM format (BEGIN PUBLIC KEY): {e}"
+        ) from e
+
+
+def _base64url_decode(data: str) -> bytes:
+    """Decode base64url-encoded string (no padding required)."""
+    padding = 4 - len(data) % 4
+    if padding != 4:
+        data += "=" * padding
+    return base64.urlsafe_b64decode(data)
+
+
+def validate_jwt(token: str, public_key: "rsa.PublicKey") -> dict:
+    """
+    Decode and verify an RS256-signed JWT.
+
+    Verifies:
+    - Token structure (3 dot-separated segments)
+    - Algorithm is RS256
+    - RSA signature using the provided public key
+    - Expiration (exp claim, if present) with clock skew tolerance
+    - Not-before (nbf claim, if present) with clock skew tolerance
+
+    Args:
+        token: Raw JWT string (without "Bearer " prefix).
+        public_key: RSA public key for signature verification.
+
+    Returns:
+        Decoded claims dict (the JWT payload).
+
+    Raises:
+        JWTValidationError: On any structural, signature, or expiration failure.
+    """
+    rsa = _require_rsa()
+
+    # Split into segments
+    parts = token.split(".")
+    if len(parts) != 3:
+        raise JWTValidationError(f"Malformed JWT: expected 3 segments, got {len(parts)}")
+
+    header_b64, payload_b64, signature_b64 = parts
+
+    # Decode header and check algorithm
+    try:
+        header = json.loads(_base64url_decode(header_b64))
+    except Exception as e:
+        raise JWTValidationError(f"Invalid JWT header: {e}") from e
+
+    alg = header.get("alg")
+    if alg != "RS256":
+        raise JWTValidationError(f"Unsupported algorithm '{alg}' — only RS256 is accepted")
+
+    # Decode signature
+    try:
+        signature = _base64url_decode(signature_b64)
+    except Exception as e:
+        raise JWTValidationError(f"Invalid JWT signature encoding: {e}") from e
+
+    # Verify RS256 signature over "<header_b64>.<payload_b64>"
+    signing_input = f"{header_b64}.{payload_b64}".encode("ascii")
+    try:
+        hash_method = rsa.verify(signing_input, signature, public_key)
+    except rsa.VerificationError:
+        raise JWTValidationError("JWT signature verification failed") from None
+
+    if hash_method != "SHA-256":
+        raise JWTValidationError(f"Unexpected hash method '{hash_method}' — expected SHA-256 for RS256")
+
+    # Decode payload
+    try:
+        claims = json.loads(_base64url_decode(payload_b64))
+    except Exception as e:
+        raise JWTValidationError(f"Invalid JWT payload: {e}") from e
+
+    # Check expiration (with clock skew tolerance for distributed systems)
+    now = time.time()
+    exp = claims.get("exp")
+    if exp is not None and now > exp + JWT_CLOCK_SKEW_SECONDS:
+        raise JWTValidationError("JWT has expired")
+
+    # Check not-before
+    nbf = claims.get("nbf")
+    if nbf is not None and now < nbf - JWT_CLOCK_SKEW_SECONDS:
+        raise JWTValidationError("JWT is not yet valid (nbf)")
+
+    return claims
+
+
+def _extract_header(event: dict, header_name: str) -> Optional[str]:
+    """Extract a header value from an API Gateway event (v1 or v2), case-insensitive."""
+    headers = event.get("headers") or {}
+    # API Gateway v2 (HTTP API) lowercases all header names
+    # API Gateway v1 (REST) preserves original case
+    for key, value in headers.items():
+        if key.lower() == header_name.lower():
+            return value
+    return None
+
+
+def require_auth(event: dict, secret_name: Optional[str] = None) -> dict:
+    """
+    Extract and validate a Bearer token from an API Gateway event.
+
+    Handles both API Gateway v1 (REST) and v2 (HTTP API) event formats.
+
+    Args:
+        event: API Gateway Lambda proxy integration event.
+        secret_name: Optional Secrets Manager secret name for the public key.
+
+    Returns:
+        Decoded JWT claims dict.
+
+    Raises:
+        AuthenticationError: If the Authorization header is missing, malformed, or token is invalid.
+    """
+    auth_header = _extract_header(event, "Authorization")
+    if not auth_header:
+        raise AuthenticationError("Missing Authorization header")
+
+    if not auth_header.startswith("Bearer "):
+        raise AuthenticationError("Authorization header must use Bearer scheme")
+
+    token = auth_header[7:]  # len("Bearer ") == 7
+    if not token:
+        raise AuthenticationError("Empty Bearer token")
+
+    try:
+        public_key = get_jwt_public_key(secret_name=secret_name)
+        return validate_jwt(token, public_key)
+    except JWTValidationError as e:
+        raise AuthenticationError(f"Authentication failed: {e}") from e
+
+
+def _normalize_path(path: str) -> str:
+    """Normalize a URL path for safe comparison.
+
+    URL-decodes, collapses duplicate slashes, ensures a single leading slash,
+    and strips any trailing slash (except for root "/").
+    """
+    path = unquote(path)
+    path = re.sub(r"/+", "/", path)
+    return "/" + path.strip("/") if path.strip("/") else "/"
+
+
+def check_auth(
+    event: dict,
+    public_paths: AbstractSet[str] = frozenset(),
+    secret_name: Optional[str] = None,
+) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
+    """Check JWT authentication on an API Gateway event, skipping public paths.
+
+    Combines path normalization, public-path bypass, JWT validation,
+    and a standard JSON:API 401 error response in one call.
+
+    Args:
+        event: API Gateway Lambda proxy integration event.
+        public_paths: Set of normalized paths that skip auth (e.g. {"/health"}).
+        secret_name: Optional Secrets Manager secret name for the public key.
+
+    Returns:
+        (claims, None) on success — claims is the decoded JWT dict,
+            or None if the path is public.
+        (None, response) on auth failure — response is a 401 dict
+            ready to return from your Lambda handler.
+    """
+    raw_path = event.get("path") or event.get("rawPath") or ""
+    if _normalize_path(raw_path) in public_paths:
+        return None, None
+
+    try:
+        claims = require_auth(event, secret_name=secret_name)
+        return claims, None
+    except AuthenticationError as e:
+        log.warning("Authentication failed: %s", e)
+        return None, {
+            "statusCode": 401,
+            "headers": {
+                "Content-Type": "application/json",
+                "Access-Control-Allow-Origin": "*",
+            },
+            "body": json.dumps({
+                "errors": [{
+                    "status": "401",
+                    "title": "Unauthorized",
+                    "detail": "Authentication required",
+                }]
+            }),
+        }

nui_lambda_shared_utils/lambda_helpers.py

@@ -0,0 +1,84 @@
+"""
+Lambda context helpers for extracting environment information.
+
+Provides standardized environment info extraction for logging and metrics context.
+"""
+
+import os
+from typing import Dict, Union
+
+
+__all__ = ["get_lambda_environment_info"]
+
+
+def get_lambda_environment_info() -> Dict[str, Union[str, bool]]:
+    """
+    Extract standard Lambda environment info.
+
+    Returns a dict with Lambda runtime information useful for logging context,
+    metric dimensions, and conditional behavior based on environment.
+
+    Detection logic:
+    - `is_local`: True if AWS_LAMBDA_RUNTIME_API is not set (local dev or tests)
+    - `environment`: Derived from ENVIRONMENT env var, falls back to "unknown"
+
+    Returns:
+        Dict with keys:
+            - environment: "prod" | "dev" | "staging" | "unknown"
+            - aws_region: AWS region (e.g., "ap-southeast-2")
+            - function_name: Lambda function name
+            - function_version: Lambda function version (e.g., "$LATEST")
+            - memory_limit: Memory limit in MB (e.g., "512")
+            - is_local: True if running outside Lambda environment
+
+    Example:
+        >>> from nui_lambda_shared_utils import get_lambda_environment_info
+        >>> env_info = get_lambda_environment_info()
+        >>> env_info
+        {
+            "environment": "prod",
+            "aws_region": "ap-southeast-2",
+            "function_name": "my-lambda",
+            "function_version": "$LATEST",
+            "memory_limit": "512",
+            "is_local": False
+        }
+
+    Usage with Powertools logger:
+        >>> logger = get_powertools_logger("my-service")
+        >>> env_info = get_lambda_environment_info()
+        >>> logger.info("Starting handler", extra=env_info)
+
+    Usage with CloudWatch metrics dimensions:
+        >>> metrics = MetricsPublisher(namespace="MyService")
+        >>> env_info = get_lambda_environment_info()
+        >>> metrics.add_dimension("Environment", env_info["environment"])
+        >>> metrics.add_dimension("FunctionName", env_info["function_name"])
+    """
+    # Detect if running in Lambda environment
+    # AWS_LAMBDA_RUNTIME_API is set by the Lambda runtime, not available locally
+    is_local = os.getenv("AWS_LAMBDA_RUNTIME_API") is None
+
+    # Environment detection
+    # Support common env var patterns: ENVIRONMENT, ENV, STAGE
+    environment = (
+        os.getenv("ENVIRONMENT")
+        or os.getenv("ENV")
+        or os.getenv("STAGE")
+        or "unknown"
+    ).lower()
+
+    # Normalize common environment names
+    if environment in ("production", "prd"):
+        environment = "prod"
+    elif environment in ("development",):
+        environment = "dev"
+
+    return {
+        "environment": environment,
+        "aws_region": os.getenv("AWS_REGION", os.getenv("AWS_DEFAULT_REGION", "")),
+        "function_name": os.getenv("AWS_LAMBDA_FUNCTION_NAME", ""),
+        "function_version": os.getenv("AWS_LAMBDA_FUNCTION_VERSION", ""),
+        "memory_limit": os.getenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", ""),
+        "is_local": is_local,
+    }