DocsGPT/docsgpt/api/devices/pairing.py

"""Pairing flow (RFC 8628 device authorization grant)."""

from __future__ import annotations

import json
import logging
import os
import secrets
import time
import uuid
from typing import Optional

from flask import jsonify, make_response, request
from sqlalchemy.exc import IntegrityError

from docsgpt.api.devices.auth import fingerprint_pubkey, hash_session_token
from docsgpt.cache import get_redis_instance
from docsgpt.core.settings import settings
from docsgpt.storage.db.repositories.devices import DevicesRepository
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
from docsgpt.storage.db.session import db_session


logger = logging.getLogger(__name__)


_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
_USER_CODE_LEN = 8  # ABCD-WXYZ (8 chars, displayed with a single dash)
_PAIRING_REDIS_PREFIX = "docsgpt:device_pairing:"
_ALLOWED_APPROVAL_MODES = {"ask", "full"}
# Attempts to find a non-colliding device name before releasing the claim.
_DEVICE_NAME_RETRIES = 5

# Atomic redeem claim: read the pairing JSON, and only if its ``status`` is
# ``pending`` flip it to ``redeemed`` and write it back (TTL preserved).
# Returns 1 if this caller won the claim, 0 otherwise (missing/not pending).
# Run server-side so two concurrent redeems of one code can't both win.
_CLAIM_PAIRING_LUA = """
local raw = redis.call('GET', KEYS[1])
if not raw then
    return 0
end
local ok, state = pcall(cjson.decode, raw)
if not ok then
    return 0
end
if state['status'] ~= 'pending' then
    return 0
end
state['status'] = 'redeemed'
local ttl = redis.call('TTL', KEYS[1])
local encoded = cjson.encode(state)
if ttl and ttl > 0 then
    redis.call('SET', KEYS[1], encoded, 'EX', ttl)
else
    redis.call('SET', KEYS[1], encoded)
end
return 1
"""


def _redis():
    return get_redis_instance()


def _generate_user_code() -> str:
    return "".join(secrets.choice(_USER_CODE_ALPHABET) for _ in range(_USER_CODE_LEN))


def _format_user_code(code: str) -> str:
    return f"{code[:4]}-{code[4:]}"


def _normalize_user_code(code: str) -> str:
    return code.replace("-", "").replace(" ", "").upper()


def _user_code_index_key(user_code: str) -> str:
    return f"docsgpt:device_pairing_code:{user_code}"


def create_pairing() -> tuple:
    """Create a new pairing code for the calling user.

    Optional body fields ``name``, ``description``, and ``approval_mode``
    are stashed in Redis alongside the pairing state and consumed by
    ``redeem_pairing`` when the device row is created. Unknown values
    fall back to: name -> CLI hostname; description -> empty; approval_mode
    -> ``ask``.
    """
    decoded = getattr(request, "decoded_token", None)
    user_id = decoded.get("sub") if isinstance(decoded, dict) else None
    if not user_id:
        return _error("authentication_required", 401)

    body = request.get_json(silent=True) or {}
    raw_name = body.get("name")
    if raw_name is not None and not isinstance(raw_name, str):
        return _error("invalid_name", 400)
    requested_name = (raw_name or "").strip() or None
    requested_description = body.get("description")
    if isinstance(requested_description, str):
        requested_description = requested_description.strip() or None
    else:
        requested_description = None
    requested_mode = body.get("approval_mode")
    if requested_mode is not None and requested_mode not in _ALLOWED_APPROVAL_MODES:
        return _error("invalid_approval_mode", 400)

    redis_client = _redis()
    if redis_client is None:
        return _error("redis_unavailable", 503)

    device_code = f"dc_{uuid.uuid4().hex}"
    user_code = _generate_user_code()
    ttl = int(settings.REMOTE_DEVICE_PAIRING_TTL_SECONDS)

    state = {
        "device_code": device_code,
        "user_code": _format_user_code(user_code),
        "user_code_raw": user_code,
        "user_id": user_id,
        "status": "pending",
        "created_at": int(time.time()),
        "requested_name": requested_name,
        "requested_description": requested_description,
        "requested_approval_mode": requested_mode,
    }
    try:
        redis_client.setex(
            _PAIRING_REDIS_PREFIX + device_code, ttl, json.dumps(state)
        )
        redis_client.setex(_user_code_index_key(user_code), ttl, device_code)
    except Exception:
        logger.exception("redis setex failed during pairing create")
        return _error("redis_unavailable", 503)

    base_url = settings.API_URL or ""
    return make_response(
        jsonify(
            {
                "device_code": device_code,
                "user_code": _format_user_code(user_code),
                "verification_uri": base_url,
                "expires_in": ttl,
                "interval": 3,
            }
        ),
        200,
    )


def get_pairing(device_code: str) -> tuple:
    """Poll pairing status (UI side)."""
    decoded = getattr(request, "decoded_token", None)
    user_id = decoded.get("sub") if isinstance(decoded, dict) else None
    if not user_id:
        return _error("authentication_required", 401)

    state = _load_pairing(device_code)
    if state is None:
        return _error("pairing_not_found", 404)
    if state.get("user_id") != user_id:
        return _error("pairing_not_found", 404)

    return make_response(
        jsonify(
            {
                "device_code": state["device_code"],
                "user_code": state.get("user_code"),
                "status": state.get("status", "pending"),
                "device_id": state.get("device_id"),
                "device_name": state.get("device_name"),
            }
        ),
        200,
    )


def delete_pairing(device_code: str) -> tuple:
    decoded = getattr(request, "decoded_token", None)
    user_id = decoded.get("sub") if isinstance(decoded, dict) else None
    if not user_id:
        return _error("authentication_required", 401)
    state = _load_pairing(device_code)
    if state is None:
        return _error("pairing_not_found", 404)
    if state.get("user_id") != user_id:
        return _error("pairing_not_found", 404)
    redis_client = _redis()
    if redis_client is None:
        return _error("redis_unavailable", 503)
    try:
        redis_client.delete(_PAIRING_REDIS_PREFIX + device_code)
        if state.get("user_code_raw"):
            redis_client.delete(_user_code_index_key(state["user_code_raw"]))
    except Exception:
        logger.exception("redis delete failed during pairing cancel")
    return make_response(jsonify({"success": True}), 200)


def redeem_pairing() -> tuple:
    """CLI submits the user_code and machine info to claim a device row.

    Returns ``device_id`` + opaque ``session_token`` (shown once). The
    token is stored only as ``token_hash`` so a Redis snapshot leak +
    DB snapshot leak still can't reconstruct it.
    """
    body = request.get_json(silent=True) or {}
    user_code_raw = _normalize_user_code(str(body.get("user_code", "")))
    if not user_code_raw or len(user_code_raw) != _USER_CODE_LEN:
        return _error("invalid_user_code", 400)
    hostname = body.get("hostname") or "unknown"
    os_name = body.get("os") or ""
    arch = body.get("arch") or ""
    cli_version = body.get("cli_version") or ""
    pubkey_b64 = body.get("machine_pubkey")
    if not pubkey_b64:
        return _error("missing_machine_pubkey", 400)

    redis_client = _redis()
    if redis_client is None:
        return _error("redis_unavailable", 503)
    try:
        device_code = redis_client.get(_user_code_index_key(user_code_raw))
        if device_code is None:
            return _error("pairing_not_found", 404)
        if isinstance(device_code, (bytes, bytearray)):
            device_code = device_code.decode("utf-8")
    except Exception:
        logger.exception("redis get failed during pairing redeem")
        return _error("redis_unavailable", 503)

    state = _load_pairing(device_code)
    if state is None:
        return _error("pairing_not_found", 404)
    if state.get("status") != "pending":
        return _error("pairing_already_redeemed", 409)

    # Atomically claim the pairing before doing any work. Two concurrent
    # redeems of the same code can both pass the check above; only the one
    # that flips ``pending`` -> ``redeemed`` here proceeds to create a device.
    claimed = _claim_pairing(redis_client, device_code)
    if claimed is None:
        return _error("redis_unavailable", 503)
    if not claimed:
        return _error("pairing_already_redeemed", 409)

    user_id = state["user_id"]

    try:
        fingerprint = fingerprint_pubkey(pubkey_b64)
    except Exception:
        return _error("invalid_machine_pubkey", 400)
    session_token = "tok_" + secrets.token_urlsafe(32)
    token_hash = hash_session_token(session_token)
    device_id = "dev_" + uuid.uuid4().hex

    # Apply UI-side overrides stashed at pairing create time, with
    # sensible fallbacks.
    requested_name = state.get("requested_name") or None
    description = state.get("requested_description") or None
    approval_mode = state.get("requested_approval_mode") or "ask"
    if approval_mode not in _ALLOWED_APPROVAL_MODES:
        approval_mode = "ask"

    name_source = requested_name or hostname
    # The claim has already been consumed, so a failed device insert would
    # strand a one-time code. ``_next_device_name`` only adds a 16-bit
    # suffix, so a ``UNIQUE (user_id, name)`` collision is possible; retry
    # with a fresh name a few times, and release the claim if every attempt
    # collides so the user can redeem the code again.
    name = None
    created = False
    for _attempt in range(_DEVICE_NAME_RETRIES):
        candidate = _next_device_name(user_id, name_source)
        try:
            with db_session() as conn:
                DevicesRepository(conn).create(
                    device_id=device_id,
                    user_id=user_id,
                    name=candidate,
                    machine_pubkey_fingerprint=fingerprint,
                    token_hash=token_hash,
                    hostname=hostname,
                    os=os_name,
                    arch=arch,
                    cli_version=cli_version,
                    approval_mode=approval_mode,
                    description=description,
                )
                _upsert_remote_device_user_tool(
                    conn, user_id=user_id, device_id=device_id,
                    device_name=candidate, description=description,
                )
            name = candidate
            created = True
            break
        except IntegrityError:
            # Name collision on this attempt; the transaction rolled back.
            # Loop to try a freshly-generated name.
            logger.warning(
                "device name collision on redeem (attempt %d); retrying",
                _attempt + 1,
            )
            continue
        except Exception:
            logger.exception("failed to create device row during redeem")
            _release_claim(redis_client, device_code, state, user_code_raw)
            return _error("internal_error", 500)

    if not created:
        logger.error(
            "device create exhausted name retries during redeem; "
            "releasing claim",
        )
        _release_claim(redis_client, device_code, state, user_code_raw)
        return _error("internal_error", 500)

    state["status"] = "redeemed"
    state["device_id"] = device_id
    state["device_name"] = name
    try:
        redis_client.setex(
            _PAIRING_REDIS_PREFIX + device_code, 300, json.dumps(state)
        )
        if state.get("user_code_raw"):
            redis_client.delete(_user_code_index_key(state["user_code_raw"]))
    except Exception:
        logger.exception("redis update failed after redeem (non-fatal)")

    return make_response(
        jsonify(
            {
                "device_id": device_id,
                "session_token": session_token,
                "name": name,
            }
        ),
        200,
    )


def _claim_pairing(redis_client, device_code: str) -> Optional[bool]:
    """Atomically transition a pairing from ``pending`` to ``redeemed``.

    Returns ``True`` if this caller won the claim, ``False`` if the pairing
    is already non-pending or gone, and ``None`` if the atomic op itself
    failed (caller should treat that as Redis-unavailable rather than
    proceeding, to avoid a non-atomic double-redeem).
    """
    key = _PAIRING_REDIS_PREFIX + device_code
    try:
        result = redis_client.eval(_CLAIM_PAIRING_LUA, 1, key)
    except Exception:
        logger.exception("redis eval failed during pairing claim")
        return None
    return bool(result)


def _release_claim(
    redis_client, device_code: str, state: dict, user_code_raw: str,
) -> None:
    """Revert a consumed claim so the one-time code is redeemable again.

    Called when the device row can't be created after the atomic claim
    already flipped the pairing to ``redeemed``. Restores ``pending`` and
    re-points the user-code index at the device code, preserving the
    remaining TTL where possible. Best-effort: failures are logged, not
    raised, since the caller is already returning an error.
    """
    key = _PAIRING_REDIS_PREFIX + device_code
    try:
        ttl = redis_client.ttl(key)
        revert = dict(state)
        revert["status"] = "pending"
        revert.pop("device_id", None)
        revert.pop("device_name", None)
        encoded = json.dumps(revert)
        if isinstance(ttl, int) and ttl > 0:
            redis_client.setex(key, ttl, encoded)
            if user_code_raw:
                redis_client.setex(
                    _user_code_index_key(user_code_raw), ttl, device_code
                )
        else:
            redis_client.set(key, encoded)
            if user_code_raw:
                redis_client.set(_user_code_index_key(user_code_raw), device_code)
    except Exception:
        logger.exception("failed to release pairing claim for %s", device_code)


def _load_pairing(device_code: str) -> Optional[dict]:
    redis_client = _redis()
    if redis_client is None:
        return None
    try:
        raw = redis_client.get(_PAIRING_REDIS_PREFIX + device_code)
    except Exception:
        return None
    if raw is None:
        return None
    if isinstance(raw, (bytes, bytearray)):
        raw = raw.decode("utf-8")
    try:
        return json.loads(raw)
    except Exception:
        return None


def _next_device_name(user_id: str, hostname: str) -> str:
    """Return a candidate name for the new device row.

    Adds a random 16-bit suffix to reduce collisions on the DB
    ``UNIQUE (user_id, name)`` constraint. Collisions are still possible,
    so ``redeem_pairing`` retries with a fresh candidate on conflict.
    """
    base = hostname or "device"
    base = base.strip() or "device"
    return f"{base}-{os.urandom(2).hex()}"


def _upsert_remote_device_user_tool(
    conn, *, user_id: str, device_id: str,
    device_name: str, description: Optional[str],
) -> None:
    """Create or update the ``user_tools`` row that fronts this device."""
    repo = UserToolsRepository(conn)
    existing = _find_remote_device_tool_row(conn, user_id, device_id)
    actions = [
        {
            "name": "run_command",
            "description": (
                f"Execute a shell command on the remote device '{device_name}'. "
                f"{description or ''}".strip()
            ),
            "active": True,
            # ``tool_executor`` consults ``RemoteDeviceTool.preview_requires_approval``
            # live for ``remote_device`` calls, so the static snapshot must
            # stay neutral. The tool's own ``execute_action`` still yields
            # the awaiting-approval pause when the live decision says so.
            "require_approval": False,
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Shell command to run.",
                        "filled_by_llm": True,
                        "value": "",
                    },
                    "working_directory": {
                        "type": "string",
                        "description": "Working directory on the remote.",
                        "filled_by_llm": True,
                        "value": "",
                    },
                    "timeout_ms": {
                        "type": "integer",
                        "description": "Timeout in milliseconds (max 600000).",
                        "filled_by_llm": True,
                        "value": "",
                    },
                },
                "required": ["command"],
            },
        }
    ]
    config = {"device_id": device_id}
    display_name = f"Remote device — {device_name}"
    if existing:
        repo.update(
            str(existing["id"]),
            user_id,
            {
                "display_name": display_name,
                "description": description or "",
                "config": config,
                "actions": actions,
                "status": True,
            },
        )
    else:
        repo.create(
            user_id=user_id,
            name="remote_device",
            display_name=display_name,
            description=description or "",
            config=config,
            actions=actions,
            status=True,
        )


def _find_remote_device_tool_row(conn, user_id: str, device_id: str):
    """Locate the user_tools row whose config.device_id matches."""
    from sqlalchemy import text
    row = conn.execute(
        text(
            """
            SELECT *
            FROM user_tools
            WHERE user_id = :user_id
              AND name    = 'remote_device'
              AND config ->> 'device_id' = :device_id
            LIMIT 1
            """
        ),
        {"user_id": user_id, "device_id": device_id},
    ).fetchone()
    if row is None:
        return None
    return dict(row._mapping)


def _error(code: str, status: int) -> tuple:
    return make_response(jsonify({"success": False, "error": code}), status)