""" Session security: concurrent-login detection, active-session management, and mid-session IP/User-Agent anomaly detection. Pen-test remediation. Everything here is inert unless SESSION_SECURITY_ENABLED=True (tools/config.py). The registry is keyed on the app's already-resolved identity (tools.gradio_platform.resolve_session_identity), so it works the same whether identity comes from an ALB "authenticate-cognito" header, Gradio-native Cognito login, or (inertly, since every "session" is then already unique) no auth at all. What this can and cannot do: - Can: track sessions per username with IP/User-Agent/login/last-seen metadata; detect concurrent logins for the same user and notify and/or soft-terminate the older session at the app layer; keep an account activity logbook; let a user list/terminate their own other sessions; detect IP/User-Agent drift mid-session and react. - Cannot: directly delete another browser's load-balancer session cookie (e.g. ALB's AWSELBAuthSessionCookie) - that store is owned by the load balancer/IdP, not this app. SESSION_SECURITY_COGNITO_GLOBAL_SIGNOUT is a best-effort, delayed (next token refresh) mitigation at the Cognito layer, not an instant kill. """ from __future__ import annotations import threading from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Any, Dict, List, Optional import boto3 from botocore.exceptions import ( BotoCoreError, ClientError, NoCredentialsError, PartialCredentialsError, ) from tools.config import ( AWS_ACCESS_KEY, AWS_REGION, AWS_SECRET_KEY, AWS_USER_POOL_ID, COGNITO_AUTH, HOST_NAME, RUN_AWS_FUNCTIONS, SAVE_LOGS_TO_CSV, SAVE_LOGS_TO_DYNAMODB, SESSION_SECURITY_ANOMALY_ACTION, SESSION_SECURITY_BIND_IP, SESSION_SECURITY_BIND_USER_AGENT, SESSION_SECURITY_COGNITO_GLOBAL_SIGNOUT, SESSION_SECURITY_DYNAMODB_TABLE_NAME, SESSION_SECURITY_ENABLED, SESSION_SECURITY_IDLE_TIMEOUT_MINUTES, SESSION_SECURITY_MODE, SESSION_SECURITY_STORE_BACKEND, ) def _now_iso() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") @dataclass class SessionRecord: username: str session_hash: str ip_address: str = "" user_agent: str = "" login_time: str = field(default_factory=_now_iso) last_seen_time: str = field(default_factory=_now_iso) status: str = "active" # active | terminated notice: Optional[str] = None def to_item(self) -> Dict[str, Any]: return { "session_hash": self.session_hash, "username": self.username, "ip_address": self.ip_address, "user_agent": self.user_agent, "login_time": self.login_time, "last_seen_time": self.last_seen_time, "status": self.status, "notice": self.notice, } @classmethod def from_item(cls, item: Dict[str, Any]) -> "SessionRecord": return cls( username=item.get("username", ""), session_hash=item.get("session_hash", ""), ip_address=item.get("ip_address", "") or "", user_agent=item.get("user_agent", "") or "", login_time=item.get("login_time") or _now_iso(), last_seen_time=item.get("last_seen_time") or _now_iso(), status=item.get("status") or "active", notice=item.get("notice") or None, ) @dataclass class HeartbeatResult: status: str notice: Optional[str] = None ### # Storage backends ### class SessionStore: """Minimal storage interface for session records.""" def get(self, session_hash: str) -> Optional[SessionRecord]: raise NotImplementedError def put(self, record: SessionRecord) -> None: raise NotImplementedError def list_for_user(self, username: str) -> List[SessionRecord]: raise NotImplementedError def delete(self, session_hash: str) -> None: raise NotImplementedError class InMemorySessionStore(SessionStore): """Default store. Per-process only: fine for a single instance/task, but concurrent logins landing on different worker processes/replicas will not see each other. Use the DynamoDB backend for multi-replica deployments.""" def __init__(self) -> None: self._lock = threading.Lock() self._records: Dict[str, SessionRecord] = {} def get(self, session_hash: str) -> Optional[SessionRecord]: with self._lock: return self._records.get(session_hash) def put(self, record: SessionRecord) -> None: with self._lock: self._records[record.session_hash] = record def list_for_user(self, username: str) -> List[SessionRecord]: with self._lock: return [r for r in self._records.values() if r.username == username] def delete(self, session_hash: str) -> None: with self._lock: self._records.pop(session_hash, None) class DynamoDBSessionStore(SessionStore): """Shares session state across multiple app replicas/tasks via DynamoDB.""" def __init__(self, table_name: str) -> None: self._table_name = table_name self._dynamodb = self._connect() self._table = self._ensure_table() @staticmethod def _connect(): try: dynamodb = boto3.resource("dynamodb", region_name=AWS_REGION) dynamodb.meta.client.list_tables() return dynamodb except Exception as exc: if RUN_AWS_FUNCTIONS and AWS_ACCESS_KEY and AWS_SECRET_KEY: return boto3.resource( "dynamodb", aws_access_key_id=AWS_ACCESS_KEY, aws_secret_access_key=AWS_SECRET_KEY, region_name=AWS_REGION, ) raise RuntimeError( f"AWS credentials for DynamoDB session store not found: {exc}" ) from exc def _ensure_table(self): try: table = self._dynamodb.Table(self._table_name) table.load() return table except ClientError as exc: if exc.response["Error"]["Code"] != "ResourceNotFoundException": raise table = self._dynamodb.create_table( TableName=self._table_name, KeySchema=[{"AttributeName": "session_hash", "KeyType": "HASH"}], AttributeDefinitions=[ {"AttributeName": "session_hash", "AttributeType": "S"} ], BillingMode="PAY_PER_REQUEST", ) table.meta.client.get_waiter("table_exists").wait( TableName=self._table_name ) return table def get(self, session_hash: str) -> Optional[SessionRecord]: response = self._table.get_item(Key={"session_hash": session_hash}) item = response.get("Item") return SessionRecord.from_item(item) if item else None def put(self, record: SessionRecord) -> None: self._table.put_item(Item=record.to_item()) def list_for_user(self, username: str) -> List[SessionRecord]: from boto3.dynamodb.conditions import Attr items: List[Dict[str, Any]] = [] scan_kwargs: Dict[str, Any] = { "FilterExpression": Attr("username").eq(username) } while True: response = self._table.scan(**scan_kwargs) items.extend(response.get("Items", [])) last_key = response.get("LastEvaluatedKey") if not last_key: break scan_kwargs["ExclusiveStartKey"] = last_key return [SessionRecord.from_item(item) for item in items] def delete(self, session_hash: str) -> None: self._table.delete_item(Key={"session_hash": session_hash}) _store: Optional[SessionStore] = None _store_lock = threading.Lock() def _get_store() -> SessionStore: global _store if _store is not None: return _store with _store_lock: if _store is None: if (SESSION_SECURITY_STORE_BACKEND or "memory").lower() == "dynamodb": try: _store = DynamoDBSessionStore(SESSION_SECURITY_DYNAMODB_TABLE_NAME) except Exception as exc: print( "Session security: could not initialise DynamoDB session store, " f"falling back to in-memory ({exc})" ) _store = InMemorySessionStore() else: _store = InMemorySessionStore() return _store def reset_store_for_tests() -> None: """Test helper: force a fresh in-memory store on next access.""" global _store _store = None ### # Account activity logbook (shared access log CSV / DynamoDB table) ### def _log_event( *, session_hash: str, username: str, event: str, ip_address: str = "", user_agent: str = "", status: str = "", details: str = "", ) -> None: """Append a security event to the existing access log (same file / DynamoDB table).""" if not SAVE_LOGS_TO_CSV and not SAVE_LOGS_TO_DYNAMODB: return try: # Local import avoids a circular import at module load # (gradio_platform does not import session_security). from tools.gradio_platform import log_platform_access log_platform_access( session_hash, HOST_NAME, username=username, event=event, ip_address=ip_address, user_agent=user_agent, status=status, details=details, ) except OSError as exc: print(f"Session security: activity log write failed ({exc})") except Exception as exc: # pragma: no cover - defensive, must never break the app print(f"Session security: could not write activity log ({exc})") ### # Cognito / Gradio best-effort session invalidation ### def _cognito_global_sign_out(username: str) -> None: """ Best-effort, delayed mitigation: invalidates the user's Cognito refresh tokens so that an ALB "authenticate-cognito" session (or any other Cognito-token consumer) is forced to re-authenticate the next time it tries to refresh. This does NOT immediately revoke an already-issued, unexpired access token, so an already-live session can remain usable for up to its access-token lifetime. It also cannot be scoped to a single sibling session - it affects every token issued for that username, including ones issued moments earlier for the very login that triggered this call. For that reason this is opt-in and off by default; do not treat it as an instant, targeted session kill. """ if ( not SESSION_SECURITY_COGNITO_GLOBAL_SIGNOUT or not AWS_USER_POOL_ID or not username ): return try: client = boto3.client("cognito-idp", region_name=AWS_REGION) client.admin_user_global_sign_out( UserPoolId=AWS_USER_POOL_ID, Username=username ) except ( ClientError, NoCredentialsError, PartialCredentialsError, BotoCoreError, ) as exc: print( "Session security: Cognito admin_user_global_sign_out failed " f"(check cognito-idp:AdminUserGlobalSignOut permission): {exc}" ) except Exception as exc: # pragma: no cover - defensive print( f"Session security: unexpected error during Cognito global sign-out: {exc}" ) _gradio_patch_lock = threading.Lock() _gradio_patch_applied = False _tracked_blocks: Any = None def _ensure_gradio_app_capture_patch() -> None: """ Best-effort monkeypatch of gradio.routes.App.configure_app so we can later reach the live App instance's ``.tokens`` dict (Gradio's own auth-cookie -> username map), in order to purge a user's Gradio-native login tokens immediately on invalidate. Only relevant when COGNITO_AUTH=True (Gradio's own login form, as opposed to ALB-header auth, where Gradio never owns the session cookie in the first place). Wrapped in try/except so that if Gradio's internals change shape in a future version, this quietly stops working instead of breaking the app. """ global _gradio_patch_applied if _gradio_patch_applied: return with _gradio_patch_lock: if _gradio_patch_applied: return try: import gradio.routes as gr_routes original_configure_app = gr_routes.App.configure_app def _patched_configure_app(self, blocks): original_configure_app(self, blocks) try: blocks._session_security_app_instance = self except Exception: # pragma: no cover - defensive pass gr_routes.App.configure_app = _patched_configure_app _gradio_patch_applied = True except Exception as exc: # pragma: no cover - defensive print( f"Session security: could not attach Gradio App-capture patch ({exc})" ) def register_gradio_blocks(blocks: Any) -> None: """Call once with the app's gr.Blocks instance to enable Gradio-native token purge.""" global _tracked_blocks if not SESSION_SECURITY_ENABLED or not COGNITO_AUTH: return _ensure_gradio_app_capture_patch() _tracked_blocks = blocks def _purge_gradio_tokens(username: str) -> None: if not COGNITO_AUTH or _tracked_blocks is None or not username: return app_instance = getattr(_tracked_blocks, "_session_security_app_instance", None) if app_instance is None: return try: tokens = getattr(app_instance, "tokens", None) if not isinstance(tokens, dict): return for token in list(tokens.keys()): if tokens.get(token) == username: del tokens[token] except Exception as exc: # pragma: no cover - defensive print(f"Session security: could not purge Gradio auth tokens ({exc})") ### # Anomaly detection ### def _detect_anomaly( record: SessionRecord, ip_address: str, user_agent: str ) -> Optional[str]: changes = [] if ( SESSION_SECURITY_BIND_IP and record.ip_address and ip_address and record.ip_address != ip_address ): changes.append(f"IP address changed from {record.ip_address} to {ip_address}") if ( SESSION_SECURITY_BIND_USER_AGENT and record.user_agent and user_agent and record.user_agent != user_agent ): changes.append("User-Agent changed") return "; ".join(changes) if changes else None def _pop_notice(store: SessionStore, record: SessionRecord) -> Optional[str]: notice = record.notice record.notice = None store.put(record) return notice ### # Public API ### def register_session( username: str, session_hash: str, ip_address: str = "", user_agent: str = "" ) -> Optional[str]: """ Register (or refresh) an active session for username/session_hash. Detects other concurrent active sessions for the same username and applies SESSION_SECURITY_MODE to them (notify and/or invalidate). Returns a pending notice for *this* session_hash, if one was left by an earlier call (e.g. this tab itself was flagged before reloading). """ if not SESSION_SECURITY_ENABLED or not username or not session_hash: return None store = _get_store() now = _now_iso() existing = store.get(session_hash) if existing is not None and existing.username == username: existing.ip_address = existing.ip_address or ip_address existing.user_agent = existing.user_agent or user_agent existing.last_seen_time = now return _pop_notice(store, existing) siblings = [ r for r in store.list_for_user(username) if r.session_hash != session_hash and r.status == "active" ] record = SessionRecord( username=username, session_hash=session_hash, ip_address=ip_address, user_agent=user_agent, login_time=now, last_seen_time=now, status="active", notice=None, ) store.put(record) _log_event( session_hash=session_hash, username=username, event="login", ip_address=ip_address, user_agent=user_agent, status="active", ) if siblings: mode = (SESSION_SECURITY_MODE or "notify").lower() do_notify = mode in ("notify", "both") do_invalidate = mode in ("invalidate", "both") details = ( f"New sign-in detected from {ip_address or 'an unknown location'} at {now}." ) for sibling in siblings: if do_invalidate: sibling.status = "terminated" sibling.notice = ( "This session was ended because a new sign-in for your account was " "detected. If this wasn't you, please review your account activity." ) event = "invalidated_by_new_login" elif do_notify: sibling.notice = ( f"A new sign-in for your account was detected from " f"{ip_address or 'a different location'}. If this wasn't you, please " "review your account activity." ) event = "notified_of_new_login" else: event = None if event: _log_event( session_hash=sibling.session_hash, username=username, event=event, ip_address=sibling.ip_address, user_agent=sibling.user_agent, status=sibling.status, details=details, ) store.put(sibling) if do_invalidate: _cognito_global_sign_out(username) _purge_gradio_tokens(username) return None def heartbeat( session_hash: str, ip_address: str = "", user_agent: str = "" ) -> HeartbeatResult: """ Periodic client poll: updates idle/last-seen tracking, checks for IP/User-Agent drift, enforces an optional idle timeout, and surfaces any pending notice (new login elsewhere, anomaly, manual termination) to the calling tab. """ if not SESSION_SECURITY_ENABLED or not session_hash: return HeartbeatResult(status="active", notice=None) store = _get_store() record = store.get(session_hash) if record is None: return HeartbeatResult(status="active", notice=None) if record.status == "terminated": return HeartbeatResult(status="terminated", notice=_pop_notice(store, record)) anomaly = _detect_anomaly(record, ip_address, user_agent) if anomaly: action = (SESSION_SECURITY_ANOMALY_ACTION or "notify").lower() details = f"Property change detected: {anomaly}" if action == "terminate": record.status = "terminated" record.notice = ( "This session was ended because we detected a change in your connection " "details (possible session hijacking). Please reload and sign in again." ) _log_event( session_hash=session_hash, username=record.username, event="anomaly_terminated", ip_address=ip_address, user_agent=user_agent, status="terminated", details=details, ) return HeartbeatResult( status="terminated", notice=_pop_notice(store, record) ) if action == "notify": record.notice = ( "We noticed a change in your connection details during this session. If " "this wasn't expected, please check your account activity." ) _log_event( session_hash=session_hash, username=record.username, event="anomaly_notified", ip_address=ip_address, user_agent=user_agent, status=record.status, details=details, ) else: _log_event( session_hash=session_hash, username=record.username, event="anomaly_logged", ip_address=ip_address, user_agent=user_agent, status=record.status, details=details, ) if ( SESSION_SECURITY_IDLE_TIMEOUT_MINUTES and SESSION_SECURITY_IDLE_TIMEOUT_MINUTES > 0 ): try: last_seen_dt = datetime.fromisoformat(record.last_seen_time) except ValueError: last_seen_dt = datetime.now(timezone.utc) idle_minutes = ( datetime.now(timezone.utc) - last_seen_dt ).total_seconds() / 60.0 if idle_minutes >= SESSION_SECURITY_IDLE_TIMEOUT_MINUTES: record.status = "terminated" record.notice = ( "This session ended due to inactivity. Please reload and sign in again." ) _log_event( session_hash=session_hash, username=record.username, event="idle_timeout", ip_address=ip_address, user_agent=user_agent, status="terminated", ) return HeartbeatResult( status="terminated", notice=_pop_notice(store, record) ) record.last_seen_time = _now_iso() record.ip_address = record.ip_address or ip_address record.user_agent = record.user_agent or user_agent return HeartbeatResult(status="active", notice=_pop_notice(store, record)) def list_sessions(username: str) -> List[SessionRecord]: """List all known sessions (active and terminated) for a username, most recent first.""" if not SESSION_SECURITY_ENABLED or not username: return [] records = _get_store().list_for_user(username) records.sort(key=lambda r: r.last_seen_time, reverse=True) return records def terminate_sessions( username: str, session_hashes: List[str], actor_session_hash: str = "" ) -> int: """ Manually terminate a user's own other sessions at the app layer. Ownership-checked: only sessions belonging to `username` are affected, and the caller's own current session_hash (`actor_session_hash`) is always skipped (use the normal logout control to end the session you are currently acting from). """ if not SESSION_SECURITY_ENABLED or not username or not session_hashes: return 0 store = _get_store() terminated = 0 for session_hash in session_hashes: if not session_hash or session_hash == actor_session_hash: continue record = store.get(session_hash) if ( record is None or record.username != username or record.status == "terminated" ): continue record.status = "terminated" record.notice = ( "This session was remotely ended by you from another tab/device." ) store.put(record) _log_event( session_hash=session_hash, username=username, event="manual_terminate", ip_address=record.ip_address, user_agent=record.user_agent, status="terminated", ) terminated += 1 return terminated