devplacepy/devplacepy/database/nodes.py

# retoor <retoor@molodetz.nl>

import hashlib
import hmac
import json
import re
import secrets
from datetime import datetime, timedelta, timezone

from sqlalchemy import text

from devplacepy.config import SECRET_KEY
from devplacepy.services.nodes.spec import DEVICE_KINDS, SPEC

from .core import _in_clause, _now_iso, db, get_table
from .soft_delete import soft_delete

LIMITS = SPEC.limits

SECRET_PREFIX = "dpn_"
SECRET_PREFIX_LENGTH = 12
SECRET_BYTES = 32
LABEL_MAX = 60
EVENT_DETAIL_MAX = 500
MEGABYTE = 1024 * 1024
BATCH_SIZE = 500
HOURS_PATTERN = re.compile(r"^([01]\d|2[0-3]):[0-5]\d$")
PLATFORMS = tuple(sorted({item.adapter for item in SPEC.destinations}))
FLAG_OPEN = "open"
FLAG_RESOLVED = "resolved"
STATUS_ACTIVE = "active"
STATUS_SUSPENDED = "suspended"
STATUS_REVOKED = "revoked"

DEFAULT_SETTINGS = {
    "enabled": False,
    "bandwidth_mb_day": LIMITS.default_donor_mb_per_day,
    "max_concurrent": LIMITS.max_concurrent_per_node,
    "active_from": "",
    "active_to": "",
    "timezone": "UTC",
    "allowed_platforms": list(PLATFORMS),
    "pause_on_battery": True,
    "public_display": False,
}


class NodeLimitError(ValueError):
    pass


def _uid() -> str:
    from devplacepy.utils import generate_uid

    return generate_uid()


def _day(moment: datetime | None = None) -> str:
    return (moment or datetime.now(timezone.utc)).date().isoformat()


def hash_ip(address: str) -> str:
    if not address:
        return ""
    return hmac.new(
        SECRET_KEY.encode(), address.encode(), hashlib.sha256
    ).hexdigest()[:32]


def _hash_secret(secret: str) -> str:
    return hashlib.sha256(secret.encode()).hexdigest()


def new_secret() -> str:
    return SECRET_PREFIX + secrets.token_urlsafe(SECRET_BYTES)


def _execute(sql: str, params: dict) -> int:
    with db:
        return db.executable.execute(text(sql), params).rowcount


def create_device(
    user_uid: str,
    label: str,
    kind: str,
    *,
    platform: str = "",
    client_version: str = "",
) -> tuple[dict, str]:
    if kind not in DEVICE_KINDS:
        raise NodeLimitError("Unknown device kind")
    secret = new_secret()
    uid = _uid()
    now = _now_iso()
    created = _execute(
        "INSERT INTO node_devices (uid, user_uid, label, kind, secret_hash, "
        "secret_prefix, platform, client_version, protocol_version, status, "
        "suspended_until, suspend_reason, country, asn, asn_class, ip_hash, "
        "last_seen_at, created_at, updated_at) "
        "SELECT :uid, :user_uid, :label, :kind, :hash, :prefix, :platform, "
        ":version, 0, 'active', '', '', '', '', 'unknown', '', '', :now, :now "
        "WHERE (SELECT COUNT(*) FROM node_devices "
        "WHERE user_uid = :user_uid AND deleted_at IS NULL) < :max_devices",
        {
            "uid": uid,
            "user_uid": user_uid,
            "label": (label or "").strip()[:LABEL_MAX],
            "kind": kind,
            "hash": _hash_secret(secret),
            "prefix": secret[:SECRET_PREFIX_LENGTH],
            "platform": (platform or "")[:40],
            "version": (client_version or "")[:40],
            "now": now,
            "max_devices": LIMITS.max_devices_per_user,
        },
    )
    if created != 1:
        raise NodeLimitError(
            f"A member can register at most {LIMITS.max_devices_per_user} devices"
        )
    return get_device(uid), secret


def get_device(device_uid: str) -> dict | None:
    return get_table("node_devices").find_one(uid=device_uid, deleted_at=None)


def _batches(values) -> list[list[str]]:
    unique = list(dict.fromkeys(value for value in values if value))
    return [unique[i : i + BATCH_SIZE] for i in range(0, len(unique), BATCH_SIZE)]


def get_devices_by_uids(device_uids) -> dict[str, dict]:
    found: dict[str, dict] = {}
    for batch in _batches(device_uids):
        placeholders, params = _in_clause(batch)
        for row in db.query(
            f"SELECT * FROM node_devices WHERE uid IN ({placeholders}) "
            "AND deleted_at IS NULL",
            **params,
        ):
            found[row["uid"]] = dict(row)
    return found


def get_enabled_by_users(user_uids) -> dict[str, bool]:
    found: dict[str, bool] = {}
    for batch in _batches(user_uids):
        placeholders, params = _in_clause(batch)
        for row in db.query(
            f"SELECT user_uid, enabled FROM node_settings WHERE user_uid IN ({placeholders}) "
            "AND deleted_at IS NULL",
            **params,
        ):
            found[row["user_uid"]] = bool(row["enabled"])
    return found


def list_devices(user_uid: str) -> list[dict]:
    return list(
        get_table("node_devices").find(
            user_uid=user_uid, deleted_at=None, order_by=["created_at"]
        )
    )


def count_devices(user_uid: str) -> int:
    return get_table("node_devices").count(user_uid=user_uid, deleted_at=None)


def find_device_by_secret(secret: str) -> dict | None:
    if not secret or not secret.startswith(SECRET_PREFIX):
        return None
    digest = _hash_secret(secret)
    for row in get_table("node_devices").find(
        secret_prefix=secret[:SECRET_PREFIX_LENGTH], deleted_at=None
    ):
        if hmac.compare_digest(row.get("secret_hash") or "", digest):
            return row
    return None


def rotate_device_secret(device_uid: str, user_uid: str) -> str | None:
    secret = new_secret()
    changed = _execute(
        "UPDATE node_devices SET secret_hash = :hash, secret_prefix = :prefix, "
        "updated_at = :now WHERE uid = :uid AND user_uid = :user_uid "
        "AND deleted_at IS NULL",
        {
            "hash": _hash_secret(secret),
            "prefix": secret[:SECRET_PREFIX_LENGTH],
            "now": _now_iso(),
            "uid": device_uid,
            "user_uid": user_uid,
        },
    )
    return secret if changed == 1 else None


def revoke_device(device_uid: str, actor_uid: str) -> int:
    get_table("node_devices").update(
        {"uid": device_uid, "status": STATUS_REVOKED, "updated_at": _now_iso()},
        ["uid"],
    )
    return soft_delete("node_devices", actor_uid, uid=device_uid)


def suspend_device(device_uid: str, reason: str, minutes: int) -> bool:
    until = (datetime.now(timezone.utc) + timedelta(minutes=minutes)).isoformat()
    return (
        _execute(
            "UPDATE node_devices SET status = 'suspended', suspended_until = :until, "
            "suspend_reason = :reason, updated_at = :now "
            "WHERE uid = :uid AND deleted_at IS NULL AND status != 'revoked'",
            {
                "until": until,
                "reason": (reason or "")[:EVENT_DETAIL_MAX],
                "now": _now_iso(),
                "uid": device_uid,
            },
        )
        == 1
    )


def unsuspend_device(device_uid: str) -> bool:
    return (
        _execute(
            "UPDATE node_devices SET status = 'active', suspended_until = '', "
            "suspend_reason = '', updated_at = :now "
            "WHERE uid = :uid AND deleted_at IS NULL AND status = 'suspended'",
            {"now": _now_iso(), "uid": device_uid},
        )
        == 1
    )


def effective_status(device: dict, moment: datetime | None = None) -> str:
    status = device.get("status") or STATUS_ACTIVE
    if status != STATUS_SUSPENDED:
        return status
    until = (device.get("suspended_until") or "").strip()
    try:
        expiry = datetime.fromisoformat(until)
    except ValueError:
        return STATUS_SUSPENDED
    if expiry.tzinfo is None:
        expiry = expiry.replace(tzinfo=timezone.utc)
    now = moment or datetime.now(timezone.utc)
    return STATUS_ACTIVE if expiry <= now else STATUS_SUSPENDED


def record_device_connection(
    device_uid: str,
    *,
    platform: str = "",
    client_version: str = "",
    protocol_version: int = 0,
    country: str = "",
    asn: str = "",
    asn_class: str = "unknown",
    ip_hash: str = "",
) -> None:
    get_table("node_devices").update(
        {
            "uid": device_uid,
            "platform": platform[:40],
            "client_version": client_version[:40],
            "protocol_version": int(protocol_version),
            "country": country[:2].upper(),
            "asn": asn[:20],
            "asn_class": asn_class,
            "ip_hash": ip_hash,
            "last_seen_at": _now_iso(),
        },
        ["uid"],
    )


def _clamped(value, low: int, high: int, default: int) -> int:
    try:
        number = int(value)
    except (TypeError, ValueError):
        return default
    return max(low, min(high, number))


def _hours(value) -> str:
    value = (value or "").strip()
    return value if HOURS_PATTERN.match(value) else ""


def _platforms(value) -> list[str]:
    if isinstance(value, str):
        try:
            value = json.loads(value)
        except ValueError:
            return list(PLATFORMS)
    if not isinstance(value, (list, tuple)):
        return list(PLATFORMS)
    return [item for item in PLATFORMS if item in set(value)]


def normalise_settings(values: dict) -> dict:
    return {
        "enabled": bool(values.get("enabled")),
        "bandwidth_mb_day": _clamped(
            values.get("bandwidth_mb_day"),
            1,
            LIMITS.max_donor_mb_per_day,
            LIMITS.default_donor_mb_per_day,
        ),
        "max_concurrent": _clamped(
            values.get("max_concurrent"),
            1,
            LIMITS.max_concurrent_per_node,
            LIMITS.max_concurrent_per_node,
        ),
        "active_from": _hours(values.get("active_from")),
        "active_to": _hours(values.get("active_to")),
        "timezone": (values.get("timezone") or "UTC").strip()[:64] or "UTC",
        "allowed_platforms": _platforms(values.get("allowed_platforms")),
        "pause_on_battery": bool(values.get("pause_on_battery")),
        "public_display": bool(values.get("public_display")),
    }


def get_node_settings(user_uid: str) -> dict:
    row = get_table("node_settings").find_one(user_uid=user_uid, deleted_at=None)
    if not row:
        return dict(DEFAULT_SETTINGS)
    return normalise_settings(row)


def save_node_settings(user_uid: str, values: dict) -> dict:
    merged = {**get_node_settings(user_uid), **values}
    clean = normalise_settings(merged)
    now = _now_iso()
    _execute(
        "INSERT INTO node_settings (uid, user_uid, enabled, bandwidth_mb_day, "
        "max_concurrent, active_from, active_to, timezone, allowed_platforms, "
        "pause_on_battery, public_display, created_at, updated_at) "
        "VALUES (:uid, :user_uid, :enabled, :mb, :concurrent, :hours_from, "
        ":hours_to, :zone, :platforms, :battery, :public, :now, :now) "
        "ON CONFLICT(user_uid) DO UPDATE SET enabled = excluded.enabled, "
        "bandwidth_mb_day = excluded.bandwidth_mb_day, "
        "max_concurrent = excluded.max_concurrent, "
        "active_from = excluded.active_from, active_to = excluded.active_to, "
        "timezone = excluded.timezone, "
        "allowed_platforms = excluded.allowed_platforms, "
        "pause_on_battery = excluded.pause_on_battery, "
        "public_display = excluded.public_display, "
        "updated_at = excluded.updated_at",
        {
            "uid": _uid(),
            "user_uid": user_uid,
            "enabled": int(clean["enabled"]),
            "mb": clean["bandwidth_mb_day"],
            "concurrent": clean["max_concurrent"],
            "hours_from": clean["active_from"],
            "hours_to": clean["active_to"],
            "zone": clean["timezone"],
            "platforms": json.dumps(clean["allowed_platforms"]),
            "battery": int(clean["pause_on_battery"]),
            "public": int(clean["public_display"]),
            "now": now,
        },
    )
    return clean


def daily_cap_bytes(user_uid: str) -> int:
    return get_node_settings(user_uid)["bandwidth_mb_day"] * MEGABYTE


def record_job(
    device_uid: str,
    user_uid: str,
    *,
    ok: bool,
    bytes_up: int = 0,
    bytes_down: int = 0,
    cache_hit: bool = False,
    cap_bytes: int | None = None,
    moment: datetime | None = None,
) -> bool:
    now = (moment or datetime.now(timezone.utc)).isoformat()
    day = _day(moment)
    up = max(0, int(bytes_up))
    down = max(0, int(bytes_down))
    cap = daily_cap_bytes(user_uid) if cap_bytes is None else int(cap_bytes)
    with db:
        db.executable.execute(
            text(
                "INSERT OR IGNORE INTO node_usage_daily (uid, device_uid, user_uid, "
                "day, requests_ok, requests_failed, bytes_up, bytes_down, "
                "cache_hits_served, integrity_failures, created_at, updated_at) "
                "VALUES (:uid, :device, :user, :day, 0, 0, 0, 0, 0, 0, :now, :now)"
            ),
            {"uid": _uid(), "device": device_uid, "user": user_uid, "day": day, "now": now},
        )
        if not ok:
            db.executable.execute(
                text(
                    "UPDATE node_usage_daily SET "
                    "requests_failed = COALESCE(requests_failed, 0) + 1, "
                    "updated_at = :now WHERE device_uid = :device AND day = :day"
                ),
                {"now": now, "device": device_uid, "day": day},
            )
            return True
        counted = db.executable.execute(
            text(
                "UPDATE node_usage_daily SET "
                "requests_ok = COALESCE(requests_ok, 0) + 1, "
                "bytes_up = COALESCE(bytes_up, 0) + :up, "
                "bytes_down = COALESCE(bytes_down, 0) + :down, "
                "cache_hits_served = COALESCE(cache_hits_served, 0) + :hit, "
                "updated_at = :now "
                "WHERE device_uid = :device AND day = :day "
                "AND COALESCE(bytes_up, 0) + COALESCE(bytes_down, 0) + :up + :down "
                "<= :cap"
            ),
            {
                "up": up,
                "down": down,
                "hit": 1 if cache_hit else 0,
                "now": now,
                "device": device_uid,
                "day": day,
                "cap": cap,
            },
        ).rowcount
        if counted != 1:
            return False
        db.executable.execute(
            text(
                "INSERT INTO node_donors (uid, user_uid, requests_ok, bytes_total, "
                "first_donated_at, last_donated_at, xp_awarded_today, xp_day, "
                "created_at, updated_at) "
                "VALUES (:uid, :user, 1, :bytes, :now, :now, 0, '', :now, :now) "
                "ON CONFLICT(user_uid) DO UPDATE SET "
                "requests_ok = COALESCE(requests_ok, 0) + 1, "
                "bytes_total = COALESCE(bytes_total, 0) + :bytes, "
                "first_donated_at = CASE WHEN COALESCE(first_donated_at, '') = '' "
                "THEN :now ELSE first_donated_at END, "
                "last_donated_at = :now, updated_at = :now"
            ),
            {"uid": _uid(), "user": user_uid, "bytes": up + down, "now": now},
        )
    return True


def record_integrity_failure(device_uid: str, user_uid: str) -> None:
    now = _now_iso()
    with db:
        db.executable.execute(
            text(
                "INSERT OR IGNORE INTO node_usage_daily (uid, device_uid, user_uid, "
                "day, requests_ok, requests_failed, bytes_up, bytes_down, "
                "cache_hits_served, integrity_failures, created_at, updated_at) "
                "VALUES (:uid, :device, :user, :day, 0, 0, 0, 0, 0, 0, :now, :now)"
            ),
            {"uid": _uid(), "device": device_uid, "user": user_uid, "day": _day(), "now": now},
        )
        db.executable.execute(
            text(
                "UPDATE node_usage_daily SET "
                "integrity_failures = COALESCE(integrity_failures, 0) + 1, "
                "updated_at = :now WHERE device_uid = :device AND day = :day"
            ),
            {"now": now, "device": device_uid, "day": _day()},
        )


def claim_xp(user_uid: str, amount: int, *, moment: datetime | None = None) -> bool:
    if amount <= 0:
        return False
    day = _day(moment)
    return (
        _execute(
            "UPDATE node_donors SET "
            "xp_awarded_today = CASE WHEN COALESCE(xp_day, '') = :day "
            "THEN COALESCE(xp_awarded_today, 0) + :amount ELSE :amount END, "
            "xp_day = :day WHERE user_uid = :user AND deleted_at IS NULL AND "
            "(CASE WHEN COALESCE(xp_day, '') = :day "
            "THEN COALESCE(xp_awarded_today, 0) ELSE 0 END) + :amount <= :cap",
            {"day": day, "amount": amount, "user": user_uid, "cap": LIMITS.xp_daily_cap},
        )
        == 1
    )


def get_donor(user_uid: str) -> dict:
    row = get_table("node_donors").find_one(user_uid=user_uid, deleted_at=None)
    return {
        "requests_ok": int((row or {}).get("requests_ok") or 0),
        "bytes_total": int((row or {}).get("bytes_total") or 0),
        "first_donated_at": (row or {}).get("first_donated_at") or "",
        "last_donated_at": (row or {}).get("last_donated_at") or "",
    }


def usage_since(user_uid: str, since_day: str) -> dict:
    rows = db.query(
        "SELECT COALESCE(SUM(requests_ok), 0) AS requests_ok, "
        "COALESCE(SUM(requests_failed), 0) AS requests_failed, "
        "COALESCE(SUM(bytes_up), 0) AS bytes_up, "
        "COALESCE(SUM(bytes_down), 0) AS bytes_down "
        "FROM node_usage_daily WHERE user_uid = :user AND day >= :since "
        "AND deleted_at IS NULL",
        user=user_uid,
        since=since_day,
    )
    return dict(next(iter(rows)))


def usage_today(device_uid: str) -> dict:
    row = get_table("node_usage_daily").find_one(
        device_uid=device_uid, day=_day(), deleted_at=None
    )
    row = row or {}
    return {
        key: int(row.get(key) or 0)
        for key in (
            "requests_ok",
            "requests_failed",
            "bytes_up",
            "bytes_down",
            "cache_hits_served",
            "integrity_failures",
        )
    }


def community_stats() -> dict:
    rows = db.query(
        "SELECT COALESCE(SUM(requests_ok), 0) AS requests, "
        "COALESCE(SUM(bytes_total), 0) AS bytes, COUNT(*) AS donors "
        "FROM node_donors WHERE deleted_at IS NULL AND COALESCE(requests_ok, 0) > 0"
    )
    row = next(iter(rows))
    return {
        "requests": int(row["requests"]),
        "bytes": int(row["bytes"]),
        "donors": int(row["donors"]),
    }


def add_event(
    kind: str, *, device_uid: str = "", user_uid: str = "", detail: str = ""
) -> None:
    get_table("node_events").insert(
        {
            "uid": _uid(),
            "device_uid": device_uid,
            "user_uid": user_uid,
            "kind": kind[:40],
            "detail": (detail or "")[:EVENT_DETAIL_MAX],
            "created_at": _now_iso(),
        }
    )


def raise_flag(kind: str, *, device_uid: str = "", user_uid: str = "", detail: str = "") -> dict:
    table = get_table("node_flags")
    existing = table.find_one(
        kind=kind, device_uid=device_uid, status=FLAG_OPEN, deleted_at=None
    )
    if existing:
        return existing
    uid = _uid()
    now = _now_iso()
    table.insert(
        {
            "uid": uid,
            "device_uid": device_uid,
            "user_uid": user_uid,
            "kind": kind[:40],
            "detail": (detail or "")[:EVENT_DETAIL_MAX],
            "status": FLAG_OPEN,
            "resolved_by": "",
            "resolved_at": "",
            "created_at": now,
            "updated_at": now,
            "deleted_at": None,
            "deleted_by": None,
        }
    )
    return table.find_one(uid=uid)


def list_flags(status: str = FLAG_OPEN) -> list[dict]:
    return list(
        get_table("node_flags").find(
            status=status, deleted_at=None, order_by=["-created_at"]
        )
    )


def resolve_flag(flag_uid: str, actor_uid: str) -> bool:
    return (
        _execute(
            "UPDATE node_flags SET status = 'resolved', resolved_by = :actor, "
            "resolved_at = :now, updated_at = :now "
            "WHERE uid = :uid AND status = 'open' AND deleted_at IS NULL",
            {"actor": actor_uid, "now": _now_iso(), "uid": flag_uid},
        )
        == 1
    )


def prune_usage_daily(*, dry_run: bool = False, moment: datetime | None = None) -> int:
    cutoff = _day((moment or datetime.now(timezone.utc)) - timedelta(days=LIMITS.usage_daily_retention_days))
    count = next(
        iter(
            db.query(
                "SELECT COUNT(*) AS c FROM node_usage_daily WHERE day < :cutoff",
                cutoff=cutoff,
            )
        )
    )["c"]
    if not dry_run and count:
        _execute("DELETE FROM node_usage_daily WHERE day < :cutoff", {"cutoff": cutoff})
    return count


def prune_events(*, dry_run: bool = False, moment: datetime | None = None) -> int:
    cutoff = (
        (moment or datetime.now(timezone.utc))
        - timedelta(days=LIMITS.event_retention_days)
    ).isoformat()
    aged = next(
        iter(
            db.query(
                "SELECT COUNT(*) AS c FROM node_events WHERE created_at < :cutoff",
                cutoff=cutoff,
            )
        )
    )["c"]
    total = get_table("node_events").count()
    excess = max(0, total - aged - LIMITS.event_max_rows)
    if not dry_run:
        if aged:
            _execute("DELETE FROM node_events WHERE created_at < :cutoff", {"cutoff": cutoff})
        if excess:
            _execute(
                "DELETE FROM node_events WHERE id IN "
                "(SELECT id FROM node_events ORDER BY created_at, id LIMIT :n)",
                {"n": excess},
            )
    return aged + excess


def prune_requests(*, dry_run: bool = False, moment: datetime | None = None) -> int:
    cutoff = (
        (moment or datetime.now(timezone.utc))
        - timedelta(seconds=LIMITS.request_retention_seconds)
    ).isoformat()
    count = next(
        iter(
            db.query(
                "SELECT COUNT(*) AS c FROM node_requests WHERE created_at < :cutoff",
                cutoff=cutoff,
            )
        )
    )["c"]
    if not dry_run and count:
        _execute("DELETE FROM node_requests WHERE created_at < :cutoff", {"cutoff": cutoff})
    return count