ForcePilot/backend/package/yuxi/channel/application/service/auth_service.py

from __future__ import annotations

import base64
import hmac

from yuxi.channel.domain.port.rate_limit_port import RateLimitPort


class AuthService:
    def __init__(
        self,
        rate_limit_port: RateLimitPort,
        *,
        token: str | None = None,
        password: str | None = None,
        max_attempts: int = 5,
        lockout_seconds: int = 300,
        allow_anonymous: bool = False,
    ) -> None:
        self._rate_limiter = rate_limit_port
        self._token = token
        self._password = password
        self._max_attempts = max_attempts
        self._lockout_seconds = lockout_seconds
        self._allow_anonymous = allow_anonymous

    def update_credentials(self, *, token: str | None = None, password: str | None = None) -> None:
        if token is not None:
            self._token = token
        if password is not None:
            self._password = password

    async def authenticate(self, auth_header: str, *, client_id: str = "unknown") -> tuple[bool, str]:
        if not self._token and not self._password:
            if self._allow_anonymous:
                return True, ""
            return False, "no credentials configured"

        locked, remaining = await self._rate_limiter.is_locked(f"channel:auth:lockout:{client_id}")
        if locked:
            return False, f"auth rate limited, retry after {remaining}s"

        if self._check_credential(auth_header):
            await self._rate_limiter.reset(f"channel:auth:attempts:{client_id}")
            return True, ""

        allowed = await self._rate_limiter.check_and_incr(
            f"channel:auth:attempts:{client_id}",
            max_attempts=self._max_attempts,
            window_seconds=self._lockout_seconds,
            lockout_seconds=self._lockout_seconds,
        )
        if not allowed:
            return False, "auth rate limited"

        return False, "auth failed"

    def _check_credential(self, auth_header: str) -> bool:
        if auth_header.startswith("Bearer "):
            if self._token:
                return _safe_compare(auth_header[7:], self._token)
        if auth_header.startswith("Basic "):
            try:
                decoded = base64.b64decode(auth_header[6:]).decode()
                if ":" in decoded and self._password:
                    _, pwd = decoded.split(":", 1)
                    return _safe_compare(pwd, self._password)
            except Exception:
                pass
        return False


def _safe_compare(a: str, b: str) -> bool:
    return hmac.compare_digest(a.encode(), b.encode())