BoltClam/tests/test_auth.py

from __future__ import annotations

import base64
import hashlib
import http.server
import json
import socket
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
from urllib.parse import parse_qs, urlparse

import pytest

from boltclam.auth import (
    AuthError,
    AuthService,
    CredentialStore,
    LoopbackAuthReceiver,
    pkce_pair,
)
from boltclam.config import AccountCache

ACCOUNT = {
    "user": {"id": "1", "email": "aaron@example.com", "name": "Aaron",
             "username": "aaron"},
    "subscription": {"tier": "premium", "label": "Premium", "status": "active"},
    "balance": {"weekly_remaining": 100, "topup": 0, "allowed": True},
    "features": {"licensed": True},
}


class _Handler(http.server.BaseHTTPRequestHandler):
    def _json(self, status: int, payload: dict) -> None:
        data = json.dumps(payload).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self) -> None:  # noqa: N802
        if self.path.startswith("/boltclam/account"):
            key = self.headers.get("X-Desk-Api-Key", "")
            if key in ("bolt_live_good", "bolt_live_new"):
                self._json(200, ACCOUNT)
            else:
                self._json(401, {"error": "invalid_credentials"})
            return
        self._json(404, {"error": "not_found"})

    def do_POST(self) -> None:  # noqa: N802
        length = int(self.headers.get("Content-Length") or 0)
        raw = self.rfile.read(length) if length else b"{}"
        if self.path == "/desk/auth/oauth/exchange":
            payload = json.loads(raw or b"{}")
            self.server.exchanges.append(payload)  # type: ignore[attr-defined]
            self._json(200, {"api_key": "bolt_live_new", "user": ACCOUNT["user"]})
            return
        if self.path == "/desk/auth/logout":
            self._json(200, {})
            return
        self._json(404, {"error": "not_found"})

    def log_message(self, fmt: str, *args) -> None:
        return


@pytest.fixture
def web_server():
    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
    server.exchanges = []  # type: ignore[attr-defined]
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        yield server
    finally:
        server.shutdown()
        server.server_close()


def _auth(tmp_path, base_url) -> AuthService:
    return AuthService(
        base_url, install_id="test-install",
        credentials=CredentialStore(tmp_path / "credentials.dat"),
        account_cache=AccountCache(tmp_path / "account.json"))


def test_pkce_pair_matches_s256():
    verifier, challenge = pkce_pair()
    expected = base64.urlsafe_b64encode(
        hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode()
    assert challenge == expected
    assert 43 <= len(verifier) <= 128


def test_credential_store_fallback_round_trip(tmp_path):
    store = CredentialStore(tmp_path / "credentials.dat")
    store.save("bolt_live_secret")
    assert store.load() == "bolt_live_secret"
    store.delete()
    assert store.load() is None


def test_key_login_restore_and_logout(tmp_path, web_server):
    base = f"http://127.0.0.1:{web_server.server_address[1]}"
    auth = _auth(tmp_path, base)
    session = auth.login_with_key("bolt_live_good")
    assert session.api_key == "bolt_live_good"
    assert session.account["user"]["email"] == "aaron@example.com"
    assert auth.stored_key() == "bolt_live_good"

    restored = auth.restore()
    assert restored.account["subscription"]["tier"] == "premium"
    assert restored.offline is False

    auth.logout()
    assert auth.stored_key() is None
    assert auth.account_cache.load() is None


def test_key_login_rejects_a_bad_key(tmp_path, web_server):
    base = f"http://127.0.0.1:{web_server.server_address[1]}"
    auth = _auth(tmp_path, base)
    with pytest.raises(AuthError):
        auth.login_with_key("bolt_live_wrong")


def test_browser_login_completes_the_loopback_pkce_handoff(tmp_path, web_server):
    base = f"http://127.0.0.1:{web_server.server_address[1]}"
    auth = _auth(tmp_path, base)
    opened: dict[str, str] = {}

    def open_url(url: str) -> bool:
        opened["url"] = url
        parsed = urlparse(url)
        query = parse_qs(parsed.query)
        redirect_uri = query["redirect_uri"][0]
        state = query["state"][0]

        def callback() -> None:
            time.sleep(0.2)
            with urllib.request.urlopen(
                    f"{redirect_uri}?code=one-time&state={state}", timeout=5) as response:
                response.read()

        threading.Thread(target=callback, daemon=True).start()
        return True

    session = auth.browser_login(open_url, timeout=10)
    assert session.api_key == "bolt_live_new"
    assert auth.stored_key() == "bolt_live_new"
    exchange = web_server.exchanges[-1]  # type: ignore[attr-defined]
    assert exchange["client"] == "boltclam"
    assert exchange["code"] == "one-time"
    assert exchange["code_verifier"]
    assert exchange["redirect_uri"].startswith("http://127.0.0.1:")
    assert "device_name=BoltClam" in opened["url"]


def test_loopback_rejects_the_wrong_state():
    receiver = LoopbackAuthReceiver("expected-state")
    receiver.start()
    try:
        try:
            urllib.request.urlopen(
                f"{receiver.redirect_uri}?code=bad&state=wrong", timeout=5).read()
            assert False, "expected a 400"
        except urllib.error.HTTPError as exc:
            assert exc.code == 400
        assert receiver.event.is_set() is False

        with urllib.request.urlopen(
                f"{receiver.redirect_uri}?code=good&state=expected-state",
                timeout=5) as response:
            assert response.status == 200
        result = receiver.wait(2)
        assert result["code"] == "good"
    finally:
        receiver.close()


def test_restore_offline_uses_the_cached_account(tmp_path):
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        port = sock.getsockname()[1]
    credential_path = tmp_path / "credentials.dat"
    CredentialStore(credential_path).save("bolt_live_cached")
    cache = AccountCache(tmp_path / "account.json")
    cache.save(ACCOUNT)
    auth = AuthService(
        f"http://127.0.0.1:{port}", install_id="offline",
        credentials=CredentialStore(credential_path), account_cache=cache)
    session = auth.restore()
    assert session.offline is True
    assert session.account["user"]["email"] == "aaron@example.com"