youtube-automate/youtube_automate/web/auth.py

"""Password hashing, session cookies and CSRF tokens.

The admin UI is publicly reachable over HTTPS, so this has to be real. The design
goal from specs.md §11 is "log in once per device, effectively never again",
which means a long-lived signed cookie rather than HTTP basic auth.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import secrets
import time

SCRYPT_N = 2**14
SCRYPT_R = 8
SCRYPT_P = 1
DKLEN = 32

SESSION_MAX_AGE = 365 * 24 * 3600  # one year
COOKIE_NAME = "yta_session"

# Login throttling: after this many consecutive failures from one address, refuse
# for LOCKOUT_SECONDS regardless of whether the password is right.
MAX_FAILURES = 5
LOCKOUT_SECONDS = 60


def _b64(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def _unb64(text: str) -> bytes:
    padding = "=" * (-len(text) % 4)
    return base64.urlsafe_b64decode(text + padding)


# --------------------------------------------------------------------------
# passwords


def hash_password(password: str, *, salt: bytes | None = None) -> str:
    salt = salt if salt is not None else secrets.token_bytes(16)
    derived = hashlib.scrypt(
        password.encode("utf-8"),
        salt=salt,
        n=SCRYPT_N,
        r=SCRYPT_R,
        p=SCRYPT_P,
        dklen=DKLEN,
    )
    return f"scrypt${SCRYPT_N}${SCRYPT_R}${SCRYPT_P}${_b64(salt)}${_b64(derived)}"


def verify_password(stored: str, password: str) -> bool:
    if not stored:
        return False
    try:
        scheme, n, r, p, salt_b64, hash_b64 = stored.split("$")
        if scheme != "scrypt":
            return False
        derived = hashlib.scrypt(
            password.encode("utf-8"),
            salt=_unb64(salt_b64),
            n=int(n),
            r=int(r),
            p=int(p),
            dklen=len(_unb64(hash_b64)),
        )
    except (ValueError, TypeError):
        return False
    return hmac.compare_digest(derived, _unb64(hash_b64))


# --------------------------------------------------------------------------
# sessions


def new_secret() -> str:
    return _b64(secrets.token_bytes(32))


def _sign(secret: str, payload: bytes) -> str:
    return _b64(hmac.new(secret.encode("utf-8"), payload, hashlib.sha256).digest())


def issue_session(secret: str, *, issued_at: float | None = None) -> str:
    payload = json.dumps(
        {"iat": int(issued_at if issued_at is not None else time.time())},
        separators=(",", ":"),
    ).encode("utf-8")
    return f"{_b64(payload)}.{_sign(secret, payload)}"


def verify_session(secret: str, token: str, *, now: float | None = None) -> bool:
    if not token or not secret:
        return False
    try:
        payload_b64, signature = token.split(".", 1)
        payload = _unb64(payload_b64)
    except (ValueError, TypeError):
        return False
    if not hmac.compare_digest(_sign(secret, payload), signature):
        return False
    try:
        issued_at = int(json.loads(payload)["iat"])
    except (ValueError, KeyError, TypeError):
        return False
    age = (now if now is not None else time.time()) - issued_at
    return 0 <= age <= SESSION_MAX_AGE


def cookie_header(token: str, *, secure: bool = True) -> str:
    parts = [
        f"{COOKIE_NAME}={token}",
        "Path=/",
        "HttpOnly",
        "SameSite=Lax",
        f"Max-Age={SESSION_MAX_AGE}",
    ]
    if secure:
        parts.insert(2, "Secure")
    return "; ".join(parts)


def clear_cookie_header() -> str:
    return f"{COOKIE_NAME}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0"


def cookie_value(header: str, name: str = COOKIE_NAME) -> str:
    """Pull one cookie out of a request's Cookie header.

    Deliberately hand-rolled rather than using http.cookies.SimpleCookie. That
    parser silently discards the remainder of the header the moment it meets a
    value it considers illegal — a JSON-ish value such as `prefs={"a":1}` is
    enough — so any cookie appearing after it becomes invisible. A browser sends
    us every cookie on the domain, including ones set by unrelated services, so
    one stray value would otherwise make a perfectly valid session vanish and
    bounce the user back to the login page with no error shown.
    """
    for part in (header or "").split(";"):
        candidate, separator, value = part.strip().partition("=")
        if separator and candidate.strip() == name:
            return value.strip().strip('"')
    return ""


# --------------------------------------------------------------------------
# CSRF


def csrf_token(secret: str, session_token: str) -> str:
    return _sign(secret, b"csrf:" + session_token.encode("utf-8"))


def verify_csrf(secret: str, session_token: str, submitted: str) -> bool:
    if not submitted:
        return False
    return hmac.compare_digest(csrf_token(secret, session_token), submitted)


# --------------------------------------------------------------------------
# throttling


class LoginThrottle:
    """In-memory consecutive-failure tracker keyed by remote address."""

    def __init__(self, max_failures: int = MAX_FAILURES, lockout: int = LOCKOUT_SECONDS):
        self.max_failures = max_failures
        self.lockout = lockout
        self._state: dict[str, tuple[int, float]] = {}

    def locked(self, key: str, *, now: float | None = None) -> bool:
        failures, last = self._state.get(key, (0, 0.0))
        if failures < self.max_failures:
            return False
        elapsed = (now if now is not None else time.time()) - last
        if elapsed >= self.lockout:
            self._state.pop(key, None)
            return False
        return True

    def record_failure(self, key: str, *, now: float | None = None) -> None:
        failures, _ = self._state.get(key, (0, 0.0))
        self._state[key] = (failures + 1, now if now is not None else time.time())

    def record_success(self, key: str) -> None:
        self._state.pop(key, None)