Commonwealth-Online-Server/server/ban_store.py

"""Persistent, fail-loud IP ban storage for Commonwealth Online servers."""

from __future__ import annotations

import ipaddress
import json
import os
import tempfile
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any


@dataclass(frozen=True)
class BanEntry:
    ip: str
    reason: str = ""
    banned_at: float = 0.0

    def to_dict(self) -> dict[str, Any]:
        return {"ip": self.ip, "reason": self.reason, "bannedAt": self.banned_at}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> BanEntry | None:
        raw_ip = str(data.get("ip", "")).strip()
        if not raw_ip:
            return None
        try:
            normalized_ip = str(ipaddress.ip_address(raw_ip))
        except ValueError:
            return None
        reason = str(data.get("reason", "") or "")[:1024]
        banned_at = data.get("bannedAt", data.get("banned_at", 0.0))
        try:
            banned_at_f = float(banned_at) if banned_at is not None else 0.0
        except (TypeError, ValueError):
            banned_at_f = 0.0
        return cls(ip=normalized_ip, reason=reason, banned_at=banned_at_f)


class BanStore:
    """Thread-safe IP ban list with atomic JSON persistence.

    Existing ban files are treated as security state. If one exists but cannot be parsed,
    startup fails instead of silently replacing the effective ban list with an empty one.
    """

    def __init__(self, path: str | Path | None = None) -> None:
        self._path = Path(path) if path is not None else None
        self._lock = threading.RLock()
        self._bans: dict[str, BanEntry] = {}
        if self._path is not None:
            self.load()

    @property
    def path(self) -> Path | None:
        return self._path

    def set_path(self, path: str | Path | None) -> None:
        with self._lock:
            self._path = Path(path) if path is not None else None

    @staticmethod
    def _normalize_ip(ip: str) -> str:
        text = str(ip).strip()
        if not text:
            raise ValueError("IP address cannot be empty.")
        try:
            return str(ipaddress.ip_address(text))
        except ValueError as error:
            raise ValueError(f"Invalid IP address: {text!r}") from error

    def load(self) -> None:
        with self._lock:
            if self._path is None or not self._path.exists():
                self._bans = {}
                return
            try:
                with self._path.open("r", encoding="utf-8") as handle:
                    data = json.load(handle)
            except (OSError, UnicodeError, json.JSONDecodeError) as error:
                raise RuntimeError(f"Could not load ban file {self._path}: {error}") from error
            if not isinstance(data, dict) or not isinstance(data.get("banned_ips", []), list):
                raise RuntimeError(f"Ban file {self._path} has an invalid schema.")

            bans: dict[str, BanEntry] = {}
            invalid_entries = 0
            for item in data.get("banned_ips", []):
                if not isinstance(item, dict):
                    invalid_entries += 1
                    continue
                entry = BanEntry.from_dict(item)
                if entry is None:
                    invalid_entries += 1
                    continue
                bans[entry.ip] = entry
            if invalid_entries:
                raise RuntimeError(
                    f"Ban file {self._path} contains {invalid_entries} invalid entr{'y' if invalid_entries == 1 else 'ies'}."
                )
            self._bans = bans

    def save(self) -> None:
        with self._lock:
            if self._path is None:
                return
            self._path.parent.mkdir(parents=True, exist_ok=True)
            payload = {
                "banned_ips": [
                    entry.to_dict()
                    for entry in sorted(self._bans.values(), key=lambda entry: entry.ip)
                ]
            }
            fd, temporary_name = tempfile.mkstemp(
                prefix=f".{self._path.name}.",
                suffix=".tmp",
                dir=str(self._path.parent),
                text=True,
            )
            temporary_path = Path(temporary_name)
            try:
                if os.name != "nt":
                    os.fchmod(fd, 0o600)
                with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as handle:
                    json.dump(payload, handle, indent=2, ensure_ascii=False, allow_nan=False)
                    handle.write("\n")
                    handle.flush()
                    os.fsync(handle.fileno())
                os.replace(temporary_path, self._path)
                if os.name != "nt":
                    os.chmod(self._path, 0o600)
            except Exception:
                try:
                    os.close(fd)
                except OSError:
                    pass
                try:
                    temporary_path.unlink(missing_ok=True)
                except OSError:
                    pass
                raise

    def is_banned(self, ip: str) -> bool:
        try:
            normalized = self._normalize_ip(ip)
        except ValueError:
            return False
        with self._lock:
            return normalized in self._bans

    def get_ban(self, ip: str) -> BanEntry | None:
        try:
            normalized = self._normalize_ip(ip)
        except ValueError:
            return None
        with self._lock:
            return self._bans.get(normalized)

    def ban_ip(self, ip: str, reason: str = "") -> BanEntry:
        normalized = self._normalize_ip(ip)
        entry = BanEntry(ip=normalized, reason=str(reason or "")[:1024], banned_at=time.time())
        with self._lock:
            self._bans[normalized] = entry
            self.save()
        return entry

    def unban_ip(self, ip: str) -> bool:
        try:
            normalized = self._normalize_ip(ip)
        except ValueError:
            return False
        with self._lock:
            if normalized not in self._bans:
                return False
            del self._bans[normalized]
            self.save()
            return True

    def list_bans(self) -> list[BanEntry]:
        with self._lock:
            return [
                BanEntry(ip=entry.ip, reason=entry.reason, banned_at=entry.banned_at)
                for entry in sorted(self._bans.values(), key=lambda entry: entry.ip)
            ]