ForcePilot/backend/package/yuxi/channel/gateway/device_auth.py

import base64
import logging
import time

from yuxi.channel.gateway.auth import GatewayAuthMode, GatewayAuthResult
from yuxi.channel.gateway.device_identity import verify_signature
from yuxi.channel.gateway.rbac import GatewayRole

logger = logging.getLogger(__name__)

DEVICE_TOKEN_PREFIX = "dv."
DEVICE_TOKEN_TIMESTAMP_TOLERANCE_SECONDS = 300


def parse_device_token(token: str) -> tuple[str, int, bytes] | None:
    """解析设备令牌，格式: dv.{base64(device_id)}.{base64(timestamp_ms)}.{base64(signature)}"""
    if not token.startswith(DEVICE_TOKEN_PREFIX):
        return None
    parts = token[len(DEVICE_TOKEN_PREFIX) :].split(".")
    if len(parts) != 3:
        return None
    try:
        def _add_padding(s: str) -> str:
            missing = (4 - len(s) % 4) % 4
            return s + "=" * missing if missing else s

        device_id = base64.urlsafe_b64decode(_add_padding(parts[0])).decode("utf-8")
        timestamp_ms = int(base64.urlsafe_b64decode(_add_padding(parts[1])).decode("utf-8"))
        signature = base64.urlsafe_b64decode(_add_padding(parts[2]))
    except (ValueError, UnicodeDecodeError, base64.binascii.Error):
        return None
    return device_id, timestamp_ms, signature


def build_device_token(device_id: str, timestamp_ms: int, signature: bytes) -> str:
    """构建设备令牌"""
    device_id_b64 = base64.urlsafe_b64encode(device_id.encode("utf-8")).rstrip(b"=").decode("ascii")
    ts_b64 = base64.urlsafe_b64encode(str(timestamp_ms).encode("utf-8")).rstrip(b"=").decode("ascii")
    sig_b64 = base64.urlsafe_b64encode(signature).rstrip(b"=").decode("ascii")
    return f"{DEVICE_TOKEN_PREFIX}{device_id_b64}.{ts_b64}.{sig_b64}"


def build_challenge(device_id: str, timestamp_ms: int) -> bytes:
    return f"{device_id}:{timestamp_ms}".encode("utf-8")


def validate_timestamp(timestamp_ms: int) -> bool:
    now_ms = int(time.time() * 1000)
    if timestamp_ms > now_ms:
        return False
    diff_ms = now_ms - timestamp_ms
    return diff_ms <= DEVICE_TOKEN_TIMESTAMP_TOLERANCE_SECONDS * 1000


async def authenticate_device(
    token: str,
    lookup_public_key,
    client_ip: str | None = None,
) -> GatewayAuthResult:
    parsed = parse_device_token(token)
    if parsed is None:
        return GatewayAuthResult(
            authenticated=False,
            mode=GatewayAuthMode.DEVICE_TOKEN,
            error="设备令牌格式无效。预期格式: dv.{deviceId}.{ts}.{sig}",
        )

    device_id, timestamp_ms, signature = parsed

    if not validate_timestamp(timestamp_ms):
        return GatewayAuthResult(
            authenticated=False,
            mode=GatewayAuthMode.DEVICE_TOKEN,
            error="设备令牌时间戳超出容忍范围（±5 分钟）。",
            metadata={"device_id": device_id},
        )

    public_key_pem = await lookup_public_key(device_id)
    if public_key_pem is None:
        return GatewayAuthResult(
            authenticated=False,
            mode=GatewayAuthMode.DEVICE_TOKEN,
            error=f"未知设备: {device_id}",
            metadata={"device_id": device_id},
        )

    challenge = build_challenge(device_id, timestamp_ms)
    if not verify_signature(public_key_pem, challenge, signature):
        return GatewayAuthResult(
            authenticated=False,
            mode=GatewayAuthMode.DEVICE_TOKEN,
            error="设备签名验证失败。",
            metadata={"device_id": device_id},
        )

    logger.info(
        "Device authenticated: device_id=%s ip=%s",
        device_id,
        client_ip or "unknown",
    )

    return GatewayAuthResult(
        authenticated=True,
        user_id=f"device:{device_id}",
        mode=GatewayAuthMode.DEVICE_TOKEN,
        metadata={"device_id": device_id},
        roles=[GatewayRole.OPERATOR],
    )