titan/device_auth.py

"""
Hardware handshake: HMAC-based device authentication.

A device proves its identity by computing an HMAC over a server-issued
challenge using the shared hardware ID as the key. The hardware ID is read
from the ``ARK_S10_HWID`` environment variable.

This implementation uses :func:`hmac.compare_digest` for constant-time
comparison to avoid timing side-channels.
"""

from __future__ import annotations

import hashlib
import hmac
import os
import secrets
from typing import Optional

from .constants import DEFAULT_HMAC_ALGORITHM

HWID_ENV_VAR = "ARK_S10_HWID"


class DeviceAuthError(Exception):
    """Raised when device authentication fails or cannot be performed."""


class HardwareHandshake:
    """Perform HMAC-based device authentication.

    Parameters
    ----------
    hwid:
        Shared hardware identifier used as the HMAC key. If ``None`` (the
        default) the value is read from the ``ARK_S10_HWID`` environment
        variable.
    algorithm:
        Hash algorithm name accepted by :mod:`hashlib`. Defaults to
        ``"sha256"``.
    """

    def __init__(
        self,
        hwid: Optional[str] = None,
        algorithm: str = DEFAULT_HMAC_ALGORITHM,
    ) -> None:
        resolved = hwid if hwid is not None else os.environ.get(HWID_ENV_VAR)
        if not resolved:
            raise DeviceAuthError(
                f"hardware ID not provided and {HWID_ENV_VAR} is not set"
            )
        if algorithm not in hashlib.algorithms_available:
            raise DeviceAuthError(f"unsupported hash algorithm: {algorithm!r}")

        self._key = resolved.encode("utf-8")
        self.algorithm = algorithm

    @staticmethod
    def generate_challenge(nbytes: int = 32) -> bytes:
        """Return a cryptographically random challenge of ``nbytes`` bytes."""
        if nbytes <= 0:
            raise ValueError("nbytes must be positive")
        return secrets.token_bytes(nbytes)

    def sign(self, challenge: bytes) -> str:
        """Compute the hex-encoded HMAC of ``challenge``."""
        if not isinstance(challenge, (bytes, bytearray)):
            raise TypeError("challenge must be bytes")
        return hmac.new(self._key, bytes(challenge), self.algorithm).hexdigest()

    def verify(self, challenge: bytes, signature: str) -> bool:
        """Return ``True`` iff ``signature`` is a valid HMAC for ``challenge``.

        Comparison is performed in constant time.
        """
        if not isinstance(signature, str):
            return False
        expected = self.sign(challenge)
        return hmac.compare_digest(expected, signature)