lockwatch 0.1.0__py3-none-any.whl

lockwatch/__init__.py

@@ -0,0 +1,47 @@
+"""
+LockWatch — API security middleware for FastAPI and Flask.
+
+Public API surface:
+
+    from lockwatch import LockWatchMiddleware, RateLimiter, JWTRotationManager
+    from lockwatch import AnomalyDetector, AuditLogger
+
+Quickstart (FastAPI):
+
+    from fastapi import FastAPI
+    from lockwatch import LockWatchMiddleware
+
+    app = FastAPI()
+    app.add_middleware(
+        LockWatchMiddleware,
+        redis_url="redis://localhost:6379",
+        rate_limit_requests=100,
+        rate_limit_window_seconds=60,
+    )
+"""
+
+from lockwatch.rate_limiter import RateLimiter, RateLimitConfig, RateLimitState
+from lockwatch.jwt_rotation import JWTRotationManager, TokenPair, JWTConfig
+from lockwatch.anomaly import AnomalyDetector, AnomalyConfig
+from lockwatch.audit import AuditLogger, AuditLogEntry, AuditQuery
+from lockwatch.middleware import LockWatchMiddleware, LockWatchFlaskMiddleware, flask_jwt_required
+
+__version__ = "0.1.0"
+__all__ = [
+    # Middleware (primary entrypoint)
+    "LockWatchMiddleware",
+    "LockWatchFlaskMiddleware",
+    "flask_jwt_required",
+    # Individual components (for custom wiring)
+    "RateLimiter",
+    "RateLimitConfig",
+    "RateLimitState",
+    "JWTRotationManager",
+    "TokenPair",
+    "JWTConfig",
+    "AnomalyDetector",
+    "AnomalyConfig",
+    "AuditLogger",
+    "AuditLogEntry",
+    "AuditQuery",
+]

lockwatch/anomaly.py

@@ -0,0 +1,161 @@
+"""
+IP burst anomaly detector.
+
+Uses a separate Redis sorted set per IP (distinct from rate limiter keys) to
+track request timestamps in a short burst window. When count exceeds threshold,
+fires async webhook alerts and records a dedup key so we don't spam on every
+subsequent request from the same IP.
+
+Alert dedup: after firing, set Redis key lockwatch:anomaly:alerted:{ip} with
+TTL=60s. Skip alert if key exists. This means at most one alert per IP per minute.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Optional
+
+import httpx
+from redis.asyncio import Redis
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AnomalyConfig:
+    """
+    Configuration for burst anomaly detection.
+
+    burst_threshold: number of requests in burst_window_seconds that triggers an alert.
+    webhook_urls: list of HTTPS endpoints to POST alert payloads to.
+    alert_cooldown_seconds: minimum seconds between alerts for the same IP.
+    """
+
+    burst_threshold: int = 50
+    burst_window_seconds: int = 10
+    webhook_urls: list[str] = field(default_factory=list)
+    alert_cooldown_seconds: int = 60
+    key_prefix: str = "lockwatch:anomaly"
+
+
+@dataclass
+class AnomalyEvent:
+    """Payload sent to configured webhook URLs when a burst is detected."""
+
+    ip: str
+    timestamp_utc: str
+    request_count_in_window: int
+    window_seconds: int
+    endpoint: str
+    method: str
+    user_agent: Optional[str] = None
+
+
+class AnomalyDetector:
+    """
+    Per-IP burst detector with async webhook alerting.
+
+    Usage:
+        detector = AnomalyDetector(redis_client, AnomalyConfig(burst_threshold=50, burst_window_seconds=10))
+        is_anomaly = await detector.check(ip="1.2.3.4", endpoint="/api/login", method="POST")
+    """
+
+    def __init__(self, redis: Redis, config: AnomalyConfig) -> None:
+        self._redis = redis
+        self._config = config
+        # Shared httpx client — reuse connections across alert dispatches
+        self._http: Optional[httpx.AsyncClient] = None
+
+    async def check(
+        self,
+        ip: str,
+        endpoint: str,
+        method: str,
+        user_agent: Optional[str] = None,
+    ) -> bool:
+        """
+        Record a request from ip and return True if a burst anomaly is detected.
+
+        Does NOT block the request — that decision is left to the middleware.
+        Alert dispatch is fire-and-forget (asyncio.create_task) so it never
+        adds latency to the response path.
+        """
+        now_ms = int(time.time() * 1000)
+        window_start_ms = now_ms - (self._config.burst_window_seconds * 1000)
+        burst_key = self._ip_redis_key(ip)
+
+        import uuid
+        pipe = self._redis.pipeline()
+        # UUID suffix ensures uniqueness within the same millisecond — same reason as rate_limiter
+        member = f"{now_ms}:{uuid.uuid4().hex}"
+        pipe.zadd(burst_key, {member: now_ms})
+        pipe.zremrangebyscore(burst_key, 0, window_start_ms)
+        pipe.zcard(burst_key)
+        pipe.expire(burst_key, self._config.burst_window_seconds)
+        results = await pipe.execute()
+
+        count: int = results[2]
+        is_burst = count >= self._config.burst_threshold
+
+        if is_burst and self._config.webhook_urls:
+            # Only fire if we haven't alerted for this IP within the cooldown window
+            alerted_key = self._alerted_redis_key(ip)
+            already_alerted = await self._redis.exists(alerted_key)
+            if not already_alerted:
+                await self._redis.set(alerted_key, "1", ex=self._config.alert_cooldown_seconds)
+                event = AnomalyEvent(
+                    ip=ip,
+                    timestamp_utc=datetime.now(timezone.utc).isoformat(),
+                    request_count_in_window=count,
+                    window_seconds=self._config.burst_window_seconds,
+                    endpoint=endpoint,
+                    method=method,
+                    user_agent=user_agent,
+                )
+                # Fire-and-forget — never let webhook delay block a response
+                asyncio.create_task(self._dispatch_alert(event))
+
+        return is_burst
+
+    async def _dispatch_alert(self, event: AnomalyEvent) -> None:
+        """
+        POST alert payload to all configured webhook URLs concurrently.
+
+        Failures are logged but swallowed — a broken webhook should never affect
+        the application. Each webhook gets a 5s timeout.
+        """
+        if self._http is None:
+            self._http = httpx.AsyncClient(timeout=5.0)
+
+        payload = {
+            "ip": event.ip,
+            "timestamp": event.timestamp_utc,
+            "request_count_in_window": event.request_count_in_window,
+            "window_seconds": event.window_seconds,
+            "endpoint": event.endpoint,
+            "method": event.method,
+            "user_agent": event.user_agent,
+        }
+
+        async def post_to(url: str) -> None:
+            try:
+                await self._http.post(url, json=payload)  # type: ignore[union-attr]
+            except Exception as exc:
+                logger.warning("LockWatch anomaly webhook failed (%s): %s", url, exc)
+
+        await asyncio.gather(*[post_to(url) for url in self._config.webhook_urls])
+
+    async def close(self) -> None:
+        """Clean up the shared HTTP client. Call on application shutdown."""
+        if self._http:
+            await self._http.aclose()
+
+    def _ip_redis_key(self, ip: str) -> str:
+        return f"{self._config.key_prefix}:burst:{ip}"
+
+    def _alerted_redis_key(self, ip: str) -> str:
+        return f"{self._config.key_prefix}:alerted:{ip}"

lockwatch/audit.py

@@ -0,0 +1,179 @@
+"""
+Postgres audit log for every request processed by LockWatch middleware.
+
+Design choices:
+  - Writes are async background tasks (fire-and-forget) so audit logging is
+    never on the critical response path. A failed write logs a warning but
+    does not affect the response.
+  - SQLAlchemy 2.0 async session with asyncpg driver — fastest available Python
+    Postgres driver; async avoids blocking the event loop.
+  - Indices on (timestamp, user_id) and (timestamp, ip) support the two most
+    common audit query patterns: "what did user X do?" and "what came from IP Y?".
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# SQLAlchemy imports are conditional — audit is an optional dependency
+try:
+    from sqlalchemy import Index, String, Boolean, Integer, Float, DateTime, select
+    from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+    from sqlalchemy.ext.asyncio import AsyncSession, AsyncEngine, create_async_engine, async_sessionmaker
+
+    class _Base(DeclarativeBase):
+        pass
+
+    class AuditLogEntry(_Base):
+        """
+        SQLAlchemy model for a single audited request.
+
+        Every request that passes through LockWatchMiddleware gets one row.
+        Columns are intentionally wide — storage is cheap, retroactive debugging
+        of security incidents is not.
+        """
+
+        __tablename__ = "lockwatch_audit_log"
+
+        id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+        # When the request was received (UTC, stored as naive datetime per SQLAlchemy convention)
+        timestamp: Mapped[datetime] = mapped_column(DateTime, nullable=False)
+        # Authenticated user identifier if available; NULL for unauthenticated requests
+        user_id: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
+        ip: Mapped[str] = mapped_column(String(45), nullable=False)  # IPv6 max = 39 chars + CIDR
+        endpoint: Mapped[str] = mapped_column(String(2048), nullable=False)
+        method: Mapped[str] = mapped_column(String(10), nullable=False)
+        status_code: Mapped[int] = mapped_column(Integer, nullable=False)
+        # True if this request was counted against a rate limit (even if not blocked)
+        rate_limit_hit: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
+        # True if anomaly detector fired for this request
+        anomaly_detected: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
+        # Response latency in milliseconds (measured from middleware entry to exit)
+        latency_ms: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
+
+        __table_args__ = (
+            # Fast lookups: "all requests from user X in time range"
+            Index("ix_audit_timestamp_user", "timestamp", "user_id"),
+            # Fast lookups: "all requests from IP Y in time range"
+            Index("ix_audit_timestamp_ip", "timestamp", "ip"),
+        )
+
+    _SQLALCHEMY_AVAILABLE = True
+
+except ImportError:
+    _SQLALCHEMY_AVAILABLE = False
+
+    # sqlalchemy not installed — provide a stub so the rest of the library
+    # can import audit.py without crashing even if [audit] extras are absent
+    class AuditLogEntry:  # type: ignore[no-redef]
+        """Stub: install lockwatch[audit] to enable Postgres audit logging."""
+        pass
+
+
+@dataclass
+class AuditQuery:
+    """Filters for querying the audit log."""
+
+    user_id: Optional[str] = None
+    ip: Optional[str] = None
+    endpoint_prefix: Optional[str] = None
+    from_time: Optional[datetime] = None
+    to_time: Optional[datetime] = None
+    rate_limit_hits_only: bool = False
+    anomalies_only: bool = False
+    limit: int = 100
+    offset: int = 0
+
+
+class AuditLogger:
+    """
+    Async audit logger backed by Postgres via SQLAlchemy.
+
+    Usage:
+        logger = AuditLogger(database_url="postgresql+asyncpg://user:pass@host/db")
+        await logger.init()  # creates table if not exists
+        await logger.log(entry)
+    """
+
+    def __init__(self, database_url: str) -> None:
+        if not _SQLALCHEMY_AVAILABLE:
+            raise ImportError(
+                "AuditLogger requires SQLAlchemy and asyncpg. "
+                "Install with: pip install lockwatch[audit]"
+            )
+        self._database_url = database_url
+        self._engine: Optional[AsyncEngine] = None
+        self._session_factory: Optional[async_sessionmaker[AsyncSession]] = None
+
+    async def init(self) -> None:
+        """
+        Create the audit log table if it doesn't exist.
+
+        Call once at application startup. Idempotent — safe to call on every startup.
+        """
+        self._engine = create_async_engine(self._database_url, echo=False)
+        self._session_factory = async_sessionmaker(self._engine, expire_on_commit=False)
+        async with self._engine.begin() as conn:
+            await conn.run_sync(_Base.metadata.create_all)
+
+    async def log(self, entry: AuditLogEntry) -> None:
+        """
+        Persist an audit log entry asynchronously.
+
+        Designed to be called via asyncio.create_task() so it never blocks the
+        response. Failures are logged at WARNING level — never re-raised.
+        """
+        if self._session_factory is None:
+            logger.warning("AuditLogger.log() called before init() — skipping")
+            return
+        try:
+            async with self._session_factory() as session:
+                session.add(entry)
+                await session.commit()
+        except Exception as exc:
+            logger.warning("LockWatch audit log write failed: %s", exc)
+
+    async def query(self, filters: AuditQuery) -> list[AuditLogEntry]:
+        """
+        Query audit log entries matching filters.
+
+        Returns at most filters.limit rows, ordered by timestamp DESC.
+        """
+        if self._session_factory is None:
+            raise RuntimeError("AuditLogger not initialized — call await logger.init() first")
+
+        stmt = select(AuditLogEntry)
+
+        if filters.user_id is not None:
+            stmt = stmt.where(AuditLogEntry.user_id == filters.user_id)
+        if filters.ip is not None:
+            stmt = stmt.where(AuditLogEntry.ip == filters.ip)
+        if filters.endpoint_prefix is not None:
+            stmt = stmt.where(AuditLogEntry.endpoint.like(f"{filters.endpoint_prefix}%"))
+        if filters.from_time is not None:
+            stmt = stmt.where(AuditLogEntry.timestamp >= filters.from_time)
+        if filters.to_time is not None:
+            stmt = stmt.where(AuditLogEntry.timestamp <= filters.to_time)
+        if filters.rate_limit_hits_only:
+            stmt = stmt.where(AuditLogEntry.rate_limit_hit == True)  # noqa: E712
+        if filters.anomalies_only:
+            stmt = stmt.where(AuditLogEntry.anomaly_detected == True)  # noqa: E712
+
+        stmt = stmt.order_by(AuditLogEntry.timestamp.desc())
+        stmt = stmt.limit(filters.limit).offset(filters.offset)
+
+        async with self._session_factory() as session:
+            result = await session.execute(stmt)
+            return list(result.scalars().all())
+
+    async def close(self) -> None:
+        """Dispose the SQLAlchemy engine. Call on application shutdown."""
+        if self._engine:
+            await self._engine.dispose()
+            self._engine = None
+            self._session_factory = None

lockwatch/jwt_rotation.py

@@ -0,0 +1,272 @@
+"""
+JWT rotation manager: access token (15min RS256) + refresh token (7-day).
+
+Security model:
+  - Refresh tokens are stored as SHA-256(token) in Redis — raw tokens never
+    persisted, so a Redis dump doesn't leak usable tokens.
+  - Rotation is atomic: blacklist old JTI before issuing new pair. If the new
+    issue fails (e.g. Redis down), the old token stays valid — fail open is
+    intentional for availability, but configurable.
+  - Access token blacklist: stored in Redis with TTL = remaining token lifetime.
+    After expiry the key auto-deletes; no unbounded blacklist growth.
+  - JTIs are UUID4 — unguessable, no sequential enumeration possible.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Any, Optional
+
+from jose import jwt, JWTError
+from redis.asyncio import Redis
+
+
+@dataclass
+class JWTConfig:
+    """
+    Configuration for the JWT rotation manager.
+
+    Either provide rsa_private_key_pem (PEM string) or a path to a PEM file
+    via rsa_private_key_path — the manager loads the file on init.
+    """
+
+    rsa_private_key_pem: Optional[str] = None
+    rsa_private_key_path: Optional[str] = None
+    issuer: str = "lockwatch"
+    audience: str = "lockwatch-api"
+    access_token_ttl_seconds: int = 900        # 15 minutes
+    refresh_token_ttl_seconds: int = 604800    # 7 days
+    key_prefix: str = "lockwatch:jwt"
+
+
+@dataclass
+class TokenPair:
+    """A newly-issued access + refresh token pair."""
+
+    access_token: str
+    refresh_token: str
+    access_expires_at: datetime
+    refresh_expires_at: datetime
+    token_type: str = "Bearer"
+
+
+class JWTRotationManager:
+    """
+    Issues RS256 JWTs with automatic rotation and Redis-backed blacklisting.
+
+    Usage:
+        manager = JWTRotationManager(redis_client, JWTConfig(rsa_private_key_pem=pem))
+        pair = await manager.issue_token_pair(subject="user:42", claims={"roles": ["admin"]})
+        # ... on refresh endpoint ...
+        new_pair = await manager.rotate(incoming_refresh_token)
+    """
+
+    def __init__(self, redis: Redis, config: JWTConfig) -> None:
+        self._redis = redis
+        self._config = config
+        self._private_key: Optional[str] = None  # loaded lazily
+
+    def _load_private_key(self) -> str:
+        if self._private_key:
+            return self._private_key
+        if self._config.rsa_private_key_pem:
+            self._private_key = self._config.rsa_private_key_pem
+        elif self._config.rsa_private_key_path:
+            with open(self._config.rsa_private_key_path) as f:
+                self._private_key = f.read()
+        else:
+            raise ValueError("JWTConfig requires either rsa_private_key_pem or rsa_private_key_path")
+        return self._private_key
+
+    def _extract_public_key(self) -> str:
+        """Derive RSA public key from the private key PEM for verification."""
+        from cryptography.hazmat.primitives.serialization import (
+            load_pem_private_key,
+            Encoding,
+            PublicFormat,
+        )
+        private_key_pem = self._load_private_key().encode()
+        private_key = load_pem_private_key(private_key_pem, password=None)
+        return private_key.public_key().public_bytes(Encoding.PEM, PublicFormat.SubjectPublicKeyInfo).decode()
+
+    async def issue_token_pair(
+        self, subject: str, claims: Optional[dict[str, Any]] = None
+    ) -> TokenPair:
+        """
+        Issue a new access + refresh token pair for a subject.
+
+        Generates a UUID4 JTI for each token so they're individually revocable.
+        Stores SHA-256(refresh_token) in Redis with refresh_token_ttl_seconds TTL.
+        """
+        private_key = self._load_private_key()
+        now = datetime.now(timezone.utc)
+
+        access_jti = str(uuid.uuid4())
+        access_expires_at = now + timedelta(seconds=self._config.access_token_ttl_seconds)
+        access_payload: dict[str, Any] = {
+            "sub": subject,
+            "iss": self._config.issuer,
+            "aud": self._config.audience,
+            "iat": int(now.timestamp()),
+            "exp": int(access_expires_at.timestamp()),
+            "jti": access_jti,
+            "type": "access",
+        }
+        if claims:
+            access_payload.update(claims)
+
+        access_token = jwt.encode(access_payload, private_key, algorithm="RS256")
+
+        refresh_jti = str(uuid.uuid4())
+        refresh_expires_at = now + timedelta(seconds=self._config.refresh_token_ttl_seconds)
+        refresh_payload: dict[str, Any] = {
+            "sub": subject,
+            "iss": self._config.issuer,
+            "aud": self._config.audience,
+            "iat": int(now.timestamp()),
+            "exp": int(refresh_expires_at.timestamp()),
+            "jti": refresh_jti,
+            "type": "refresh",
+        }
+
+        refresh_token = jwt.encode(refresh_payload, private_key, algorithm="RS256")
+
+        # Store hash so Redis never holds a usable token — hash lookup on rotate()
+        token_hash = self._hash_token(refresh_token)
+        redis_key = self._refresh_redis_key(token_hash)
+        await self._redis.set(redis_key, refresh_jti, ex=self._config.refresh_token_ttl_seconds)
+
+        return TokenPair(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            access_expires_at=access_expires_at,
+            refresh_expires_at=refresh_expires_at,
+        )
+
+    async def rotate(self, refresh_token: str) -> TokenPair:
+        """
+        Validate a refresh token, atomically delete its hash, and issue a new pair.
+
+        Uses WATCH/MULTI/EXEC to eliminate the race condition where two concurrent
+        callers could both pass the existence check before either deletes the key.
+        If another client deletes the key between WATCH and EXEC, WatchError is
+        raised and we retry — on retry the GET returns None → raises ValueError.
+        """
+        public_key = self._extract_public_key()
+
+        try:
+            claims = jwt.decode(
+                refresh_token,
+                public_key,
+                algorithms=["RS256"],
+                audience=self._config.audience,
+                issuer=self._config.issuer,
+            )
+        except JWTError as e:
+            raise ValueError(f"Invalid refresh token: {e}") from e
+
+        if claims.get("type") != "refresh":
+            raise ValueError("Token is not a refresh token")
+
+        jti = claims.get("jti")
+        subject = claims.get("sub")
+        if not jti or not subject:
+            raise ValueError("Refresh token missing required claims")
+
+        token_hash = self._hash_token(refresh_token)
+        refresh_key = self._refresh_redis_key(token_hash)
+        blacklist_key = self._blacklist_redis_key(jti)
+
+        from redis.exceptions import WatchError
+
+        async with self._redis.pipeline() as pipe:
+            while True:
+                try:
+                    await pipe.watch(refresh_key)
+                    # pipe is now in immediate-execution mode (not buffered)
+                    stored_jti = await pipe.get(refresh_key)
+                    if not stored_jti:
+                        await pipe.unwatch()
+                        raise ValueError("Refresh token has already been used or revoked")
+
+                    is_blacklisted = await pipe.exists(blacklist_key)
+                    if is_blacklisted:
+                        await pipe.unwatch()
+                        raise ValueError("Refresh token JTI is blacklisted")
+
+                    pipe.multi()          # switch to buffered mode
+                    pipe.delete(refresh_key)
+                    await pipe.execute()  # atomic; raises WatchError if key was changed
+                    break
+                except WatchError:
+                    continue  # concurrent rotation: retry sees empty key → ValueError
+
+        return await self.issue_token_pair(subject=subject)
+
+    async def blacklist_access_token(self, token: str) -> None:
+        """
+        Revoke an access token before its natural expiry.
+
+        Computes remaining TTL from token claims so the blacklist key auto-expires
+        exactly when the token would have expired naturally — no orphaned keys.
+        """
+        public_key = self._extract_public_key()
+
+        try:
+            claims = jwt.decode(
+                token,
+                public_key,
+                algorithms=["RS256"],
+                audience=self._config.audience,
+                issuer=self._config.issuer,
+            )
+        except JWTError as e:
+            raise ValueError(f"Cannot blacklist invalid token: {e}") from e
+
+        jti = claims.get("jti")
+        exp = claims.get("exp")
+        if not jti or not exp:
+            raise ValueError("Token missing jti or exp claim")
+
+        remaining_ttl = int(exp - datetime.now(timezone.utc).timestamp())
+        if remaining_ttl > 0:
+            await self._redis.set(self._blacklist_redis_key(jti), "1", ex=remaining_ttl)
+
+    async def verify_access_token(self, token: str) -> dict[str, Any]:
+        """
+        Verify signature, expiry, issuer, audience, and blacklist status.
+
+        Returns decoded claims dict on success. Raises jose.JWTError on failure.
+        Blacklist check is an O(1) Redis GET — fast enough for every request.
+        """
+        public_key = self._extract_public_key()
+
+        claims = jwt.decode(
+            token,
+            public_key,
+            algorithms=["RS256"],
+            audience=self._config.audience,
+            issuer=self._config.issuer,
+        )
+
+        if claims.get("type") != "access":
+            raise JWTError("Token is not an access token")
+
+        jti = claims.get("jti")
+        if jti and await self._redis.exists(self._blacklist_redis_key(jti)):
+            raise JWTError("Token has been revoked")
+
+        return claims
+
+    def _hash_token(self, token: str) -> str:
+        # SHA-256 hex digest — deterministic, non-reversible
+        return hashlib.sha256(token.encode()).hexdigest()
+
+    def _refresh_redis_key(self, token_hash: str) -> str:
+        return f"{self._config.key_prefix}:refresh:{token_hash}"
+
+    def _blacklist_redis_key(self, jti: str) -> str:
+        return f"{self._config.key_prefix}:blacklist:{jti}"