aps-agent/server/auth/licenses.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import time
import uuid
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import Any

from fastapi import Request
from sqlalchemy import select

from server.auth.context import IdentityContext
from server.auth.providers import AuthError
from server.db.database import get_session
from server.db.models import LicenseActivation, LicenseCode


LICENSE_DURATIONS: dict[str, int | None] = {
    "hour": 60 * 60,
    "day": 24 * 60 * 60,
    "week": 7 * 24 * 60 * 60,
    "month": 30 * 24 * 60 * 60,
    "year": 365 * 24 * 60 * 60,
    "permanent": None,
}

LICENSE_LABELS = {
    "hour": "小时授权",
    "day": "天授权",
    "week": "周授权",
    "month": "月授权",
    "year": "年授权",
    "permanent": "永久授权",
    "custom": "自定义授权",
}

DEFAULT_MOCK_CODES: dict[str, str] = {
    "APS-HOUR-DEMO": "hour",
    "APS-DAY-DEMO": "day",
    "APS-WEEK-DEMO": "week",
    "APS-MONTH-DEMO": "month",
    "APS-YEAR-DEMO": "year",
    "APS-PERP-DEMO": "permanent",
}

DESKTOP_COOKIE = "aps_desktop_session"
WEB_LICENSE_COOKIE = "aps_license_session"


def is_desktop_request(request: Request) -> bool:
    return (request.headers.get("x-aps-client") or "").strip().lower() == "desktop"


def license_clients() -> frozenset[str]:
    """授权码机制作用的客户端集合（APS_LICENSE_CLIENTS，默认仅 desktop，兼容既有部署）。"""
    raw = (os.environ.get("APS_LICENSE_CLIENTS") or "desktop").strip().lower()
    return frozenset(item.strip() for item in raw.split(",") if item.strip())


def use_license_auth(request: Request) -> bool:
    """判定当前请求是否应走授权码认证。

    - 桌面客户端：始终走授权（与既有行为一致）。
    - Web 客户端：仅当 APS_LICENSE_CLIENTS 包含 web、授权提供方为 codes（正式授权码模式）、
      且未配置 JMS 认证（JMS 环境优先，避免与 MES 集成冲突）时走授权码。
    """
    if is_desktop_request(request):
        return True
    if "web" not in license_clients():
        return False
    if (os.environ.get("APS_LICENSE_PROVIDER") or "").strip().lower() != "codes":
        return False
    from server.auth.providers import get_auth_provider

    return get_auth_provider().mode != "jms"


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


def _b64decode(raw: str) -> bytes:
    return base64.urlsafe_b64decode(raw + "=" * (-len(raw) % 4))


def _sha256(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def _device_user_id(device_hash: str) -> int:
    # Keep the generated owner ID inside signed BIGINT while remaining stable per installation.
    return int(device_hash[:15], 16)


def _utc_text(epoch: int) -> str:
    return datetime.fromtimestamp(epoch, tz=timezone.utc).isoformat()


class LicenseProvider(ABC):
    mode = "unconfigured"

    @abstractmethod
    async def activate(self, request: Request, payload: dict[str, Any]) -> tuple[IdentityContext, str]:
        raise NotImplementedError

    @abstractmethod
    async def authenticate(self, request: Request) -> IdentityContext:
        raise NotImplementedError

    async def refresh(self, request: Request) -> tuple[IdentityContext, str]:
        raise AuthError("LICENSE_NOT_CONFIGURED", "授权刷新接口尚未配置", 503)

    async def logout(self, request: Request) -> None:
        return None

    def describe(self, identity: IdentityContext) -> dict[str, Any] | None:
        if identity.auth_kind != "license" or not identity.license_type:
            return None
        return {
            "type": identity.license_type,
            "label": LICENSE_LABELS.get(identity.license_type, identity.license_type),
            "activatedAt": identity.license_activated_at,
            "expiresAt": identity.license_expires_at,
            "permanent": identity.license_expires_at is None,
        }


class UnconfiguredLicenseProvider(LicenseProvider):
    async def activate(self, request: Request, payload: dict[str, Any]) -> tuple[IdentityContext, str]:
        raise AuthError("LICENSE_NOT_CONFIGURED", "客户端授权接口尚未配置", 503)

    async def authenticate(self, request: Request) -> IdentityContext:
        raise AuthError("LICENSE_NOT_CONFIGURED", "客户端授权接口尚未配置", 503)


class MockLicenseProvider(LicenseProvider):
    """Development provider; production remains fail-closed unless explicitly configured."""

    mode = "mock"

    def __init__(self) -> None:
        secret = (os.environ.get("APS_MOCK_LICENSE_SECRET") or "").strip()
        tenant_uuid = (
            os.environ.get("APS_MOCK_LICENSE_TENANT_UUID")
            or os.environ.get("APS_MOCK_TENANT_UUID")
            or ""
        ).strip()
        missing = []
        if len(secret) < 32:
            missing.append("APS_MOCK_LICENSE_SECRET（至少 32 个字符）")
        if not tenant_uuid:
            missing.append("APS_MOCK_LICENSE_TENANT_UUID")
        if missing:
            raise AuthError(
                "LICENSE_CONFIG_INVALID",
                f"开发授权配置缺失：{', '.join(missing)}",
                503,
            )
        self.secret = secret.encode()
        self.tenant_uuid = tenant_uuid

    def _catalog(self) -> dict[str, str]:
        configured = (os.environ.get("APS_MOCK_LICENSE_CODES") or "").strip()
        if not configured:
            return dict(DEFAULT_MOCK_CODES)
        try:
            data = json.loads(configured)
        except json.JSONDecodeError as exc:
            raise AuthError("LICENSE_CONFIG_INVALID", "开发授权码配置无效", 503) from exc
        result: dict[str, str] = {}
        if isinstance(data, dict):
            for code, grant in data.items():
                license_type = grant.get("type") if isinstance(grant, dict) else grant
                if isinstance(code, str) and license_type in LICENSE_DURATIONS:
                    result[code.strip().upper()] = str(license_type)
        return result

    @staticmethod
    def _device_id(request: Request, payload: dict[str, Any] | None = None) -> str:
        device_id = (request.headers.get("x-aps-device-id") or "").strip()
        if not device_id and payload:
            device_id = str(payload.get("deviceId") or "").strip()
        if len(device_id) < 8 or len(device_id) > 256:
            raise AuthError("DEVICE_ID_REQUIRED", "无法识别当前客户端设备", 400)
        return device_id

    @staticmethod
    def _require_desktop(request: Request) -> None:
        if not is_desktop_request(request):
            raise AuthError("DESKTOP_CLIENT_REQUIRED", "授权码仅可用于桌面客户端", 403)

    def _issue(self, identity: IdentityContext, device_hash: str) -> str:
        payload = {**identity.to_dict(), "device_hash": device_hash}
        raw = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
        body = _b64encode(raw)
        signature = _b64encode(hmac.new(self.secret, body.encode("ascii"), hashlib.sha256).digest())
        return f"{body}.{signature}"

    def _identity(self, row: LicenseActivation) -> IdentityContext:
        token_expiry = row.expires_at or 4_102_444_800  # 2100-01-01; DB validation still supports revocation.
        return IdentityContext(
            user_id=row.user_id,
            username=f"desktop-{row.device_hash[:8]}",
            fullname="本机授权用户",
            tenant_uuid=row.tenant_uuid,
            roles=("desktop", "planner"),
            expires_at=token_expiry,
            auth_kind="license",
            license_type=row.license_type,
            license_activated_at=row.activated_at,
            license_expires_at=row.expires_at,
            activation_id=row.id,
        )

    def _verify_token(self, token: str, device_hash: str) -> IdentityContext:
        try:
            body, signature = token.split(".", 1)
            expected = _b64encode(hmac.new(self.secret, body.encode("ascii"), hashlib.sha256).digest())
            if not hmac.compare_digest(signature, expected):
                raise ValueError("signature")
            payload = json.loads(_b64decode(body))
            if payload.get("device_hash") != device_hash:
                raise AuthError("LICENSE_DEVICE_MISMATCH", "授权与当前设备不匹配", 401)
            activation_id = str(payload.get("activation_id") or "")
            if not activation_id:
                raise ValueError("activation")
        except AuthError:
            raise
        except Exception as exc:
            raise AuthError("LICENSE_INVALID", "客户端授权凭证无效", 401) from exc

        now = int(time.time())
        with get_session() as session:
            row = session.get(LicenseActivation, activation_id)
            if row is None or row.deleted != 0 or row.device_hash != device_hash:
                raise AuthError("LICENSE_INVALID", "客户端授权凭证无效", 401)
            if row.status != "active" or row.revoked_at is not None:
                raise AuthError("LICENSE_REVOKED", "客户端授权已被停用", 401)
            if row.expires_at is not None and row.expires_at <= now:
                row.status = "expired"
                session.commit()
                raise AuthError("LICENSE_EXPIRED", "客户端授权已到期", 401)
            if now - row.last_seen_at >= 3600:
                row.last_seen_at = now
                session.commit()
            return self._identity(row)

    async def activate(self, request: Request, payload: dict[str, Any]) -> tuple[IdentityContext, str]:
        self._require_desktop(request)
        device_id = self._device_id(request, payload)
        code = str(payload.get("code") or "").strip().upper()
        if not code:
            raise AuthError("LICENSE_CODE_REQUIRED", "请输入授权码", 400)
        license_type = self._catalog().get(code)
        if not license_type:
            raise AuthError("LICENSE_CODE_INVALID", "授权码无效", 401)

        now = int(time.time())
        device_hash = _sha256(device_id)
        # Demo codes are reusable across developer installations. Each resulting
        # activation is still device-bound; the production provider owns seat limits.
        code_hash = _sha256(f"mock:{code}:{device_hash}")
        with get_session() as session:
            row = session.scalar(select(LicenseActivation).where(
                LicenseActivation.code_hash == code_hash,
                LicenseActivation.deleted == 0,
            ))
            if row is not None:
                if row.device_hash != device_hash:
                    raise AuthError("LICENSE_DEVICE_MISMATCH", "授权与当前设备不匹配", 409)
                if row.status != "active" or row.revoked_at is not None:
                    raise AuthError("LICENSE_REVOKED", "客户端授权已被停用", 401)
                if row.expires_at is not None and row.expires_at <= now:
                    row.status = "expired"
                    session.commit()
                    raise AuthError("LICENSE_EXPIRED", "授权码已到期", 401)
            else:
                seconds = LICENSE_DURATIONS[license_type]
                row = LicenseActivation(
                    id=uuid.uuid4().hex,
                    code_hash=code_hash,
                    code_hint=f"{code[:4]}...{code[-4:]}",
                    device_hash=device_hash,
                    user_id=_device_user_id(device_hash),
                    license_type=license_type,
                    status="active",
                    activated_at=now,
                    expires_at=(now + seconds) if seconds is not None else None,
                    last_seen_at=now,
                    revoked_at=None,
                    tenant_uuid=self.tenant_uuid,
                    creator_id=None,
                    updater_id=None,
                    updated_at=_utc_text(now),
                    deleted=0,
                    created_at=_utc_text(now),
                )
                session.add(row)
                session.commit()
                session.refresh(row)
            identity = self._identity(row)
        return identity, self._issue(identity, device_hash)

    async def authenticate(self, request: Request) -> IdentityContext:
        self._require_desktop(request)
        device_hash = _sha256(self._device_id(request))
        token = request.cookies.get(DESKTOP_COOKIE) or ""
        if not token:
            raise AuthError("LICENSE_REQUIRED", "请先使用授权码激活客户端", 401)
        return self._verify_token(token, device_hash)

    async def refresh(self, request: Request) -> tuple[IdentityContext, str]:
        identity = await self.authenticate(request)
        device_hash = _sha256(self._device_id(request))
        return identity, self._issue(identity, device_hash)


ADMIN_LICENSE_CODE = "aps@Z@Agent@Admin"
ADMIN_LICENSE_ROLES = ("system", "admin", "planner", "approver", "auditor", "scheduler", "license")
DEFAULT_CODE_ROLES = ("license", "planner")


class CodeLicenseProvider(LicenseProvider):
    """正式授权码提供方：授权码即用户身份。

    每个授权码唯一对应一个稳定的 user_id（由授权码哈希派生），智能体会话、项目、
    主数据等全部按 tenant_uuid + user_id 归属到该授权码；不同授权码之间数据天然隔离，
    同一授权码可在任意浏览器/设备激活并找回同一份数据。
    内置超级管理员授权码 aps@Z@Agent@Admin（永久授权 + 全功能角色），
    其余授权码经 APS_LICENSE_CODES（JSON）发放。
    """

    mode = "codes"

    def __init__(self) -> None:
        secret = (os.environ.get("APS_LICENSE_SECRET") or "").strip()
        if len(secret) < 32:
            raise AuthError(
                "LICENSE_CONFIG_INVALID",
                "授权配置缺失：APS_LICENSE_SECRET（至少 32 个字符）",
                503,
            )
        self.secret = secret.encode()
        tenant = (os.environ.get("APS_LICENSE_TENANT_UUID") or "platform").strip()
        self.tenant_uuid = tenant or "platform"

    def _catalog(self) -> dict[str, dict[str, Any]]:
        catalog: dict[str, dict[str, Any]] = {
            ADMIN_LICENSE_CODE: {
                "type": "permanent",
                "roles": ADMIN_LICENSE_ROLES,
                "name": "aps-admin",
            },
        }
        configured = (os.environ.get("APS_LICENSE_CODES") or "").strip()
        if not configured:
            return catalog
        try:
            data = json.loads(configured)
        except json.JSONDecodeError as exc:
            raise AuthError("LICENSE_CONFIG_INVALID", "授权码配置无效（APS_LICENSE_CODES 须为 JSON）", 503) from exc
        if isinstance(data, dict):
            for code, grant in data.items():
                if not isinstance(code, str) or not code.strip():
                    continue
                if isinstance(grant, dict):
                    license_type = str(grant.get("type") or "")
                    roles = grant.get("roles")
                    name = grant.get("name")
                else:
                    license_type, roles, name = str(grant), None, None
                if license_type not in LICENSE_DURATIONS:
                    continue
                entry: dict[str, Any] = {"type": license_type}
                if isinstance(roles, (list, tuple)) and roles:
                    entry["roles"] = tuple(str(role) for role in roles)
                else:
                    entry["roles"] = DEFAULT_CODE_ROLES
                if isinstance(name, str) and name.strip():
                    entry["name"] = name.strip()[:64]
                catalog[code.strip()] = entry
        # 超管在设置面板发放的授权码（DB 登记，覆盖同名环境配置）
        catalog.update(self._db_codes())
        return catalog

    def _db_codes(self) -> dict[str, dict[str, Any]]:
        with get_session() as session:
            rows = session.scalars(select(LicenseCode).where(
                LicenseCode.tenant_uuid == self.tenant_uuid,
                LicenseCode.status == "active",
                LicenseCode.deleted == 0,
            )).all()
        result: dict[str, dict[str, Any]] = {}
        for row in rows:
            if row.license_type not in LICENSE_DURATIONS and not row.expires_at:
                continue
            db_entry: dict[str, Any] = {"type": row.license_type, "roles": DEFAULT_CODE_ROLES}
            if row.expires_at:
                db_entry["expiresAt"] = int(row.expires_at)
            if row.name:
                db_entry["name"] = row.name
            result[row.code] = db_entry
        return result

    @staticmethod
    def _scope(request: Request) -> str:
        scope = "desktop" if is_desktop_request(request) else "web"
        if scope not in license_clients():
            raise AuthError("LICENSE_CLIENT_DISABLED", "授权码在当前部署的该客户端不可用", 403)
        return scope

    @staticmethod
    def _code_user_id(tenant_uuid: str, code: str) -> int:
        # 授权码即身份：同一授权码恒定派生同一 user_id（保持在有符号 BIGINT 范围内）。
        digest = _sha256(f"aps-license:{tenant_uuid}:{code}")
        return int(digest[:15], 16)

    def _issue(self, identity: IdentityContext) -> str:
        payload = {**identity.to_dict(), "issuer": "codes"}
        raw = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
        body = _b64encode(raw)
        signature = _b64encode(hmac.new(self.secret, body.encode("ascii"), hashlib.sha256).digest())
        return f"{body}.{signature}"

    def _verify_token(self, token: str) -> IdentityContext:
        try:
            body, signature = token.split(".", 1)
            expected = _b64encode(hmac.new(self.secret, body.encode("ascii"), hashlib.sha256).digest())
            if not hmac.compare_digest(signature, expected):
                raise ValueError("signature")
            payload = json.loads(_b64decode(body))
            activation_id = str(payload.get("activation_id") or "")
            if not activation_id:
                raise ValueError("activation")
        except Exception as exc:
            raise AuthError("LICENSE_INVALID", "授权凭证无效", 401) from exc

        now = int(time.time())
        with get_session() as session:
            row = session.get(LicenseActivation, activation_id)
            if row is None or row.deleted != 0:
                raise AuthError("LICENSE_INVALID", "授权凭证无效", 401)
            if row.status != "active" or row.revoked_at is not None:
                raise AuthError("LICENSE_REVOKED", "授权已被停用", 401)
            if row.expires_at is not None and row.expires_at <= now:
                row.status = "expired"
                session.commit()
                raise AuthError("LICENSE_EXPIRED", "授权码已到期，请联系发放方续期", 401)
            code_row = session.scalar(select(LicenseCode).where(
                LicenseCode.code_hash == row.code_hash,
                LicenseCode.deleted == 0,
            ))
            if code_row is not None and code_row.status != "active":
                raise AuthError("LICENSE_REVOKED", "授权码已被停用，请联系管理员", 401)
            if now - row.last_seen_at >= 3600:
                row.last_seen_at = now
                session.commit()
        roles = payload.get("roles")
        return IdentityContext(
            user_id=row.user_id,
            username=str(payload.get("username") or "license-user"),
            fullname=str(payload.get("fullname") or "授权用户"),
            tenant_uuid=row.tenant_uuid,
            roles=tuple(str(role) for role in roles) if isinstance(roles, list) and roles else DEFAULT_CODE_ROLES,
            expires_at=row.expires_at or 4_102_444_800,  # 2100-01-01；DB 校验仍支持停用/到期
            auth_kind="license",
            license_type=row.license_type,
            license_activated_at=row.activated_at,
            license_expires_at=row.expires_at,
            activation_id=row.id,
        )

    async def activate(self, request: Request, payload: dict[str, Any]) -> tuple[IdentityContext, str]:
        self._scope(request)
        code = str(payload.get("code") or "").strip()
        if not code:
            raise AuthError("LICENSE_CODE_REQUIRED", "请输入授权码", 400)
        entry = self._catalog().get(code)
        if entry is None:
            raise AuthError("LICENSE_CODE_INVALID", "授权码无效", 401)

        now = int(time.time())
        code_hash = _sha256(f"codes:{self.tenant_uuid}:{code}")
        with get_session() as session:
            row = session.scalar(select(LicenseActivation).where(
                LicenseActivation.code_hash == code_hash,
                LicenseActivation.deleted == 0,
            ))
            if row is not None:
                if row.revoked_at is not None:
                    raise AuthError("LICENSE_REVOKED", "授权已被停用", 401)
                if row.expires_at is not None and row.expires_at <= now:
                    if row.status != "expired":
                        row.status = "expired"
                        session.commit()
                    raise AuthError("LICENSE_EXPIRED", "授权码已到期，请联系发放方续期", 401)
                if row.status != "active":
                    raise AuthError("LICENSE_REVOKED", "授权已被停用", 401)
                row.last_seen_at = now
                session.commit()
                session.refresh(row)
            else:
                seconds = LICENSE_DURATIONS.get(entry["type"])
                absolute_expiry = entry.get("expiresAt")
                if absolute_expiry:
                    expires_at = int(absolute_expiry)
                elif seconds is not None:
                    expires_at = now + seconds
                else:
                    expires_at = None
                if expires_at is not None and expires_at <= now:
                    raise AuthError("LICENSE_EXPIRED", "授权码已到期，请联系发放方续期", 401)
                row = LicenseActivation(
                    id=uuid.uuid4().hex,
                    code_hash=code_hash,
                    code_hint=f"{code[:2]}***{code[-2:]}" if len(code) > 4 else "***",
                    device_hash=_sha256(f"codes-shared:{self.tenant_uuid}"),
                    user_id=self._code_user_id(self.tenant_uuid, code),
                    license_type=entry["type"],
                    status="active",
                    activated_at=now,
                    expires_at=expires_at,
                    last_seen_at=now,
                    revoked_at=None,
                    tenant_uuid=self.tenant_uuid,
                    creator_id=None,
                    updater_id=None,
                    updated_at=_utc_text(now),
                    deleted=0,
                    created_at=_utc_text(now),
                )
                session.add(row)
                session.commit()
                session.refresh(row)

        identity = IdentityContext(
            user_id=row.user_id,
            username=str(entry.get("name") or f"lic-{row.code_hint}"),
            fullname="超级管理员" if code == ADMIN_LICENSE_CODE else str(entry.get("name") or "授权用户"),
            tenant_uuid=row.tenant_uuid,
            roles=tuple(entry.get("roles") or DEFAULT_CODE_ROLES),
            expires_at=row.expires_at or 4_102_444_800,
            auth_kind="license",
            license_type=row.license_type,
            license_activated_at=row.activated_at,
            license_expires_at=row.expires_at,
            activation_id=row.id,
        )
        return identity, self._issue(identity)

    async def authenticate(self, request: Request) -> IdentityContext:
        scope = self._scope(request)
        cookie = DESKTOP_COOKIE if scope == "desktop" else WEB_LICENSE_COOKIE
        token = request.cookies.get(cookie) or ""
        if not token:
            message = "请先使用授权码激活客户端" if scope == "desktop" else "请提供授权码激活后使用"
            raise AuthError("LICENSE_REQUIRED", message, 401)
        return self._verify_token(token)

    async def refresh(self, request: Request) -> tuple[IdentityContext, str]:
        identity = await self.authenticate(request)
        return identity, self._issue(identity)


_provider: LicenseProvider | None = None
_provider_mode: str | None = None


def get_license_provider() -> LicenseProvider:
    global _provider, _provider_mode
    mode = (os.environ.get("APS_LICENSE_PROVIDER") or "unconfigured").strip().lower()
    if _provider is None or _provider_mode != mode:
        if mode == "mock":
            _provider = MockLicenseProvider()
        elif mode == "codes":
            _provider = CodeLicenseProvider()
        else:
            _provider = UnconfiguredLicenseProvider()
        _provider_mode = mode
    return _provider