cmhub/apps/licensing/services.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import timedelta

from django.conf import settings
from django.db import IntegrityError, transaction
from django.utils import timezone

from apps.licensing.models import ClientDevice, DeviceBindingAudit, DeviceSession


class DeviceRegistrationError(Exception):
    def __init__(self, code: str, message: str):
        self.code = code
        self.message = message
        super().__init__(message)


class DeviceSessionValidationError(DeviceRegistrationError):
    pass


@dataclass(frozen=True)
class DeviceRegistrationResult:
    device: ClientDevice
    session: DeviceSession
    session_token: str
    created: bool


def _session_expiry(now):
    return now + timedelta(seconds=max(60, settings.DEVICE_SESSION_TTL_SECONDS))


def _find_or_create_device(*, user, product_code, device_id_version, device_id, public_key, platform, client_version, now):
    device_fingerprint = ClientDevice.fingerprint_device_id(device_id_version, device_id)
    public_key_fingerprint = ClientDevice.fingerprint_public_key(public_key)
    lookup = {
        "user": user,
        "product_code": product_code,
        "device_fingerprint": device_fingerprint,
    }

    device = ClientDevice.objects.select_for_update().filter(**lookup).first()
    created = False
    if device is None:
        try:
            with transaction.atomic():
                device = ClientDevice.objects.create(
                    **lookup,
                    device_id_version=device_id_version,
                    public_key_fingerprint=public_key_fingerprint,
                    platform=platform,
                    client_version=client_version,
                    last_seen_at=now,
                )
                created = True
        except IntegrityError:
            device = ClientDevice.objects.select_for_update().get(**lookup)

    if device.status != ClientDevice.Status.ACTIVE:
        raise DeviceRegistrationError("device_revoked", "设备已被吊销")
    if not device.public_key_fingerprint == public_key_fingerprint:
        raise DeviceRegistrationError("device_identity_mismatch", "设备身份校验失败")

    updates = []
    if device.platform != platform:
        device.platform = platform
        updates.append("platform")
    if device.client_version != client_version:
        device.client_version = client_version
        updates.append("client_version")
    update_after = timedelta(seconds=max(60, settings.DEVICE_ACTIVITY_UPDATE_SECONDS))
    if now - device.last_seen_at >= update_after:
        device.last_seen_at = now
        updates.append("last_seen_at")
    if updates:
        device.save(update_fields=[*updates, "updated_at"])
    return device, created


@transaction.atomic
def register_device(*, user, api_key, product_code, device_id_version, device_id, public_key, platform, client_version):
    now = timezone.now()
    device, created = _find_or_create_device(
        user=user,
        product_code=product_code,
        device_id_version=device_id_version,
        device_id=device_id,
        public_key=public_key,
        platform=platform,
        client_version=client_version,
        now=now,
    )

    # Registration rotates the short-lived session while keeping one device row.
    DeviceSession.objects.filter(
        device=device,
        revoked_at__isnull=True,
        expires_at__gt=now,
    ).update(revoked_at=now)
    raw_token = DeviceSession.generate_plaintext_token()
    session = DeviceSession.objects.create(
        device=device,
        token_hash=DeviceSession.hash_token(raw_token),
        expires_at=_session_expiry(now),
    )
    if created:
        DeviceBindingAudit.objects.create(
            user=user,
            device=device,
            api_key=api_key,
            action=DeviceBindingAudit.Action.REGISTERED,
            client_version=client_version,
        )
    DeviceBindingAudit.objects.create(
        user=user,
        device=device,
        api_key=api_key,
        action=DeviceBindingAudit.Action.SESSION_ISSUED,
        client_version=client_version,
    )
    return DeviceRegistrationResult(
        device=device,
        session=session,
        session_token=raw_token,
        created=created,
    )


@transaction.atomic
def record_device_heartbeat(session: DeviceSession, *, now=None) -> bool:
    now = now or timezone.now()
    session = (
        DeviceSession.objects.select_for_update()
        .select_related("device", "device__user")
        .get(pk=session.pk)
    )
    if not session.is_active_at(now) or session.device.status != ClientDevice.Status.ACTIVE:
        raise DeviceRegistrationError("device_revoked", "设备会话不可用")

    update_after = timedelta(seconds=max(60, settings.DEVICE_ACTIVITY_UPDATE_SECONDS))
    last_seen_at = session.device.last_seen_at
    if now - last_seen_at < update_after:
        return False

    ClientDevice.objects.filter(pk=session.device_id).update(last_seen_at=now)
    DeviceSession.objects.filter(pk=session.pk).update(last_used_at=now)
    DeviceBindingAudit.objects.create(
        user=session.device.user,
        device=session.device,
        action=DeviceBindingAudit.Action.HEARTBEAT,
        client_version=session.device.client_version,
    )
    return True


def resolve_optional_device_session(*, user, raw_token: str) -> DeviceSession | None:
    raw_token = str(raw_token or "").strip()
    if not raw_token:
        return None

    try:
        session = DeviceSession.objects.select_related("device", "device__user").get(
            token_hash=DeviceSession.hash_token(raw_token)
        )
    except DeviceSession.DoesNotExist as exc:
        raise DeviceSessionValidationError(
            "device_session_invalid",
            "设备会话无效或已过期",
        ) from exc

    if not session.matches_token(raw_token) or not session.is_active_at():
        raise DeviceSessionValidationError("device_session_invalid", "设备会话无效或已过期")
    if session.device.status != ClientDevice.Status.ACTIVE:
        raise DeviceSessionValidationError("device_revoked", "设备已被吊销")
    if session.device.user_id != user.id:
        raise DeviceSessionValidationError("device_mismatch", "设备会话不属于当前账号")
    return session