security_monitor.py

"""Connection security monitor for posh1's public, key-gated Space.

Scope, stated honestly: this operates at the APPLICATION layer only.
There is no infrastructure-level firewall/WAF/IP-table access available
to a Docker-SDK HF Space -- HF's own reverse proxy terminates TCP/HTTP
before this code ever runs. What this module actually does:

  - logs every /ws/session connection attempt (outcome, source IP if
    derivable, timestamp) to an append-only JSONL file on the durable
    bucket volume (/data -- survives container restarts, per the
    persistent-storage fix)
  - auto-blocks, at the application layer, any source IP that racks up
    repeated INVALID access-key attempts within a short window -- the
    practical equivalent of "banning" available here: such a source
    never has the real key, so blocking it costs zero legitimate access
  - flags (does not block) a *valid*-key source connecting at an
    abnormal rate, since that could mean the key leaked and is being
    reused -- a human call, not an auto-block, because it's still using
    real credentials

IP attribution caveat, stated plainly: behind HF's proxy,
`websocket.client.host` may show the proxy's own address rather than the
real client. This checks X-Forwarded-For first and falls back to
websocket.client.host, but on some proxy configurations neither is the
true origin -- this is a best-effort signal, not forensic-grade
attribution, and the log says so per entry.
"""
from __future__ import annotations

import json
import os
import time
from collections import defaultdict, deque
from pathlib import Path
from typing import Optional

FAILURE_THRESHOLD = 5           # invalid-key attempts...
FAILURE_WINDOW_S = 300          # ...within this many seconds...
BLOCK_DURATION_S = 1800         # ...triggers a block for this long.

RATE_FLAG_THRESHOLD = 10        # valid-key connects...
RATE_FLAG_WINDOW_S = 60         # ...within this many seconds -> flag (not block)

LOG_PATH = Path(os.environ.get("POSH_SECURITY_LOG", "/data/security_log.jsonl"))


class SecurityMonitor:
    def __init__(self, log_path: Path | None = None):
        self.log_path = log_path or LOG_PATH
        try:
            self.log_path.parent.mkdir(parents=True, exist_ok=True)
        except Exception:
            pass  # degrade to in-memory only if /data isn't writable in this context
        self._failures: dict[str, deque] = defaultdict(deque)
        self._valid_connects: dict[str, deque] = defaultdict(deque)
        self._blocked_until: dict[str, float] = {}

    # -- IP attribution (best-effort, documented limitation above) ----------
    @staticmethod
    def client_ip(websocket) -> str:
        xff = websocket.headers.get("x-forwarded-for")
        if xff:
            return xff.split(",")[0].strip()
        if websocket.client:
            return websocket.client.host
        return "unknown"

    # -- core decisions -------------------------------------------------------
    def is_blocked(self, ip: str) -> bool:
        until = self._blocked_until.get(ip)
        if until is None:
            return False
        if time.time() >= until:
            del self._blocked_until[ip]
            return False
        return True

    def record_attempt(self, ip: str, valid: bool, reason: str = "") -> Optional[str]:
        """Log the attempt; return a warning string if this attempt triggered
        a new block or rate flag, else None."""
        now = time.time()
        entry = {"t": now, "ip": ip, "valid": valid, "reason": reason,
                  "ip_attribution": "x-forwarded-for-or-direct, best-effort"}
        self._append(entry)

        if valid:
            dq = self._valid_connects[ip]
            dq.append(now)
            while dq and now - dq[0] > RATE_FLAG_WINDOW_S:
                dq.popleft()
            if len(dq) >= RATE_FLAG_THRESHOLD:
                warning = f"RATE FLAG: {ip} made {len(dq)} valid-key connections in {RATE_FLAG_WINDOW_S}s -- possible key leak, human review recommended"
                self._append({"t": now, "ip": ip, "event": "rate_flag", "count": len(dq)})
                return warning
            return None

        dq = self._failures[ip]
        dq.append(now)
        while dq and now - dq[0] > FAILURE_WINDOW_S:
            dq.popleft()
        if len(dq) >= FAILURE_THRESHOLD:
            self._blocked_until[ip] = now + BLOCK_DURATION_S
            warning = f"AUTO-BLOCK: {ip} after {len(dq)} invalid-key attempts in {FAILURE_WINDOW_S}s -- blocked {BLOCK_DURATION_S}s"
            self._append({"t": now, "ip": ip, "event": "auto_block", "until": now + BLOCK_DURATION_S})
            return warning
        return None

    def _append(self, entry: dict) -> None:
        try:
            with open(self.log_path, "a") as f:
                f.write(json.dumps(entry, separators=(",", ":")) + "\n")
        except Exception:
            pass  # never let logging failure take down a real session

    # -- stats sheet ----------------------------------------------------------
    def summary(self, tail_n: int = 2000) -> dict:
        """The stats sheet: read back the durable log and summarize it.
        This is the thing to check FIRST when diagnosing any Space anomaly."""
        if not self.log_path.exists():
            return {"total_entries": 0, "note": "no security log yet"}
        lines = self.log_path.read_text().splitlines()[-tail_n:]
        entries = []
        for line in lines:
            try:
                entries.append(json.loads(line))
            except json.JSONDecodeError:
                continue
        valid = sum(1 for e in entries if e.get("valid") is True)
        invalid = sum(1 for e in entries if e.get("valid") is False)
        blocks = [e for e in entries if e.get("event") == "auto_block"]
        rate_flags = [e for e in entries if e.get("event") == "rate_flag"]
        by_ip_invalid: dict[str, int] = defaultdict(int)
        for e in entries:
            if e.get("valid") is False:
                by_ip_invalid[e.get("ip", "unknown")] += 1
        now = time.time()
        currently_blocked = [ip for ip, until in self._blocked_until.items() if until > now]
        return {
            "total_entries": len(entries),
            "valid_connections": valid,
            "invalid_attempts": invalid,
            "auto_blocks_ever": len(blocks),
            "currently_blocked_ips": currently_blocked,
            "rate_flags_ever": len(rate_flags),
            "top_invalid_sources": sorted(by_ip_invalid.items(), key=lambda kv: -kv[1])[:10],
            "attribution_caveat": "IPs are X-Forwarded-For or direct socket, best-effort behind HF's proxy",
        }


_monitor: SecurityMonitor | None = None


def get_monitor() -> SecurityMonitor:
    global _monitor
    if _monitor is None:
        _monitor = SecurityMonitor()
    return _monitor