Nexus/server/application/services/auth_service.py

"""Nexus — Auth Service (Login / TOTP / JWT / brute-force protection)
Application layer — orchestrates Admin Repository + LoginAttempt + TOTP + JWT.
"""

import logging
import hashlib
import hmac
import secrets
import struct
import time
import datetime
from datetime import timezone
from typing import Optional

from server.domain.models import Admin, LoginAttempt, AuditLog
from server.domain.repositories import AdminRepository, LoginAttemptRepository, AuditLogRepository

logger = logging.getLogger("nexus.auth_service")

# Brute-force protection thresholds
MAX_LOGIN_FAILURES = 5  # Lock out after 5 failures within 15 minutes
LOCKOUT_MINUTES = 15

# JWT config
JWT_ACCESS_TOKEN_EXPIRE_MINUTES = 30
JWT_REFRESH_TOKEN_EXPIRE_DAYS = 7


class AuthService:
    """Business logic for authentication, TOTP, JWT, and session management"""

    def __init__(
        self,
        admin_repo: AdminRepository,
        attempt_repo: LoginAttemptRepository,
        audit_repo: AuditLogRepository,
    ):
        self.admin_repo = admin_repo
        self.attempt_repo = attempt_repo
        self.audit_repo = audit_repo

    async def login(self, username: str, password: str, ip_address: str, totp_code: Optional[str] = None) -> dict:
        """Authenticate admin user with password + optional TOTP → return JWT tokens

        Returns:
            {"success": True, "access_token": "...", "refresh_token": "...", "admin": {...}}
            or {"success": False, "reason": "..."}
        """
        # Check brute-force lockout
        failures = await self.attempt_repo.count_recent_failures(username, ip_address, LOCKOUT_MINUTES)
        if failures >= MAX_LOGIN_FAILURES:
            await self._audit("login_locked", "admin", 0, f"Account locked: {username} from {ip_address}")
            return {"success": False, "reason": "account_locked", "message": "登录尝试过多，请15分钟后重试"}

        # Find admin
        admin = await self.admin_repo.get_by_username(username)
        if not admin:
            await self._record_attempt(username, ip_address, False)
            return {"success": False, "reason": "invalid_credentials", "message": "用户名或密码错误"}

        if not admin.is_active:
            await self._record_attempt(username, ip_address, False)
            return {"success": False, "reason": "account_disabled", "message": "账户已被禁用"}

        # Verify password
        if not self._verify_password(password, admin.password_hash):
            await self._record_attempt(username, ip_address, False)
            return {"success": False, "reason": "invalid_credentials", "message": "用户名或密码错误"}

        # Verify TOTP (if enabled)
        if admin.totp_enabled:
            if not totp_code:
                return {"success": False, "reason": "totp_required", "message": "需要TOTP验证码"}
            if not self._verify_totp(totp_code, admin.totp_secret):
                await self._record_attempt(username, ip_address, False)
                return {"success": False, "reason": "invalid_totp", "message": "TOTP验证码错误"}

        # Success — generate JWT tokens
        access_token = self._create_access_token(admin)
        refresh_token = self._create_refresh_token()

        # Store refresh token in DB
        admin.jwt_refresh_token = refresh_token
        admin.jwt_token_expires = datetime.datetime.now(timezone.utc) + datetime.timedelta(days=JWT_REFRESH_TOKEN_EXPIRE_DAYS)
        admin.last_login = datetime.datetime.now(timezone.utc)
        await self.admin_repo.update(admin)

        await self._record_attempt(username, ip_address, True)
        await self._audit("login_success", "admin", admin.id, f"Login: {username} from {ip_address}")

        return {
            "success": True,
            "access_token": access_token,
            "refresh_token": refresh_token,
            "token_type": "bearer",
            "expires_in": JWT_ACCESS_TOKEN_EXPIRE_MINUTES * 60,
            "admin": {
                "id": admin.id,
                "username": admin.username,
                "email": admin.email,
                "totp_enabled": admin.totp_enabled,
            },
        }

    async def refresh_token(self, refresh_token: str) -> dict:
        """Exchange refresh token for new access token"""
        admin = await self.admin_repo.get_by_refresh_token(refresh_token)
        if not admin:
            return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}

        if admin.jwt_token_expires and admin.jwt_token_expires < datetime.datetime.now(timezone.utc):
            return {"success": False, "reason": "token_expired", "message": "刷新令牌已过期"}

        access_token = self._create_access_token(admin)
        new_refresh = self._create_refresh_token()

        admin.jwt_refresh_token = new_refresh
        admin.jwt_token_expires = datetime.datetime.now(timezone.utc) + datetime.timedelta(days=JWT_REFRESH_TOKEN_EXPIRE_DAYS)
        await self.admin_repo.update(admin)

        return {
            "success": True,
            "access_token": access_token,
            "refresh_token": new_refresh,
            "token_type": "bearer",
            "expires_in": JWT_ACCESS_TOKEN_EXPIRE_MINUTES * 60,
        }

    async def verify_access_token(self, token: str) -> Optional[Admin]:
        """Verify JWT access token → return Admin or None"""
        payload = self._decode_access_token(token)
        if not payload:
            return None

        admin_id = payload.get("sub")
        if not admin_id:
            return None

        try:
            admin_id = int(admin_id)
        except (ValueError, TypeError):
            return None

        admin = await self.admin_repo.get_by_id(admin_id)
        if not admin or not admin.is_active:
            return None

        return admin

    async def logout(self, admin_id: int) -> dict:
        """Invalidate refresh token by admin ID"""
        admin = await self.admin_repo.get_by_id(admin_id)
        if admin:
            admin.jwt_refresh_token = None
            admin.jwt_token_expires = None
            await self.admin_repo.update(admin)
        return {"success": True, "message": "已登出"}

    async def logout_by_token(self, refresh_token: str) -> dict:
        """Invalidate refresh token by token value (used by frontend logout)"""
        admin = await self.admin_repo.get_by_refresh_token(refresh_token)
        if admin:
            admin.jwt_refresh_token = None
            admin.jwt_token_expires = None
            await self.admin_repo.update(admin)
            await self._audit("logout", "admin", admin.id, f"Logout: {admin.username}")
        return {"success": True, "message": "已登出"}

    async def setup_totp(self, admin_id: int) -> dict:
        """Generate TOTP secret for an admin user (before enabling)"""
        import base64
        secret = base64.b32encode(secrets.token_bytes(20)).decode()
        admin = await self.admin_repo.get_by_id(admin_id)
        if not admin:
            return {"success": False, "reason": "admin_not_found"}

        admin.totp_secret = secret
        # Don't enable yet — user must verify once first
        await self.admin_repo.update(admin)

        # Generate provisioning URI for authenticator apps
        from server.config import settings
        issuer = settings.SYSTEM_NAME
        uri = f"otpauth://totp/{issuer}:{admin.username}?secret={secret}&issuer={issuer}"

        await self._audit("setup_totp", "admin", admin_id, f"TOTP setup initiated for {admin.username}")
        return {"success": True, "secret": secret, "uri": uri}

    async def enable_totp(self, admin_id: int, totp_code: str) -> dict:
        """Enable TOTP after user verifies with a code"""
        admin = await self.admin_repo.get_by_id(admin_id)
        if not admin or not admin.totp_secret:
            return {"success": False, "reason": "no_secret"}

        if not self._verify_totp(totp_code, admin.totp_secret):
            return {"success": False, "reason": "invalid_totp", "message": "验证码错误，TOTP未启用"}

        admin.totp_enabled = True
        await self.admin_repo.update(admin)
        await self._audit("enable_totp", "admin", admin_id, f"TOTP enabled for {admin.username}")
        return {"success": True, "message": "TOTP已启用"}

    async def disable_totp(self, admin_id: int) -> dict:
        """Disable TOTP for an admin user"""
        admin = await self.admin_repo.get_by_id(admin_id)
        if not admin:
            return {"success": False, "reason": "admin_not_found"}

        admin.totp_enabled = False
        admin.totp_secret = None
        await self.admin_repo.update(admin)
        await self._audit("disable_totp", "admin", admin_id, f"TOTP disabled for {admin.username}")
        return {"success": True, "message": "TOTP已禁用"}

    # ── JWT Token Helpers ──

    def _create_access_token(self, admin: Admin) -> str:
        """Create JWT access token with admin claims"""
        import jwt
        from server.config import settings

        now = datetime.datetime.now(timezone.utc)
        payload = {
            "sub": str(admin.id),
            "username": admin.username,
            "iat": now,
            "exp": now + datetime.timedelta(minutes=JWT_ACCESS_TOKEN_EXPIRE_MINUTES),
        }
        return jwt.encode(payload, settings.SECRET_KEY, algorithm="HS256")

    def _create_refresh_token(self) -> str:
        """Create opaque refresh token (random, stored in DB)"""
        return secrets.token_urlsafe(48)

    def _decode_access_token(self, token: str) -> Optional[dict]:
        """Decode and verify JWT access token"""
        try:
            import jwt
            from server.config import settings
            return jwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"])
        except jwt.ExpiredSignatureError:
            return None
        except jwt.InvalidTokenError:
            return None

    # ── Password & TOTP Helpers ──

    def _verify_password(self, password: str, password_hash: str) -> bool:
        """Verify password against stored hash (bcrypt-compatible)"""
        try:
            import bcrypt
            return bcrypt.checkpw(password.encode(), password_hash.encode())
        except ImportError:
            # Fallback: SHA256 comparison (for initial setup before bcrypt installed)
            return hashlib.sha256(password.encode()).hexdigest() == password_hash

    def _verify_totp(self, code: str, secret: str) -> bool:
        """Verify TOTP code using RFC 6238 (HOTP with time counter)"""
        import base64
        key = base64.b32decode(secret, casefold=True)
        # Current time step (30-second intervals)
        counter = int(time.time()) // 30

        # Check current and ±1 window (allows 30s clock drift)
        for offset in range(-1, 2):
            c = counter + offset
            # HOTP algorithm
            msg = struct.pack(">Q", c)
            h = hmac.new(key, msg, hashlib.sha1).digest()
            offset_byte = h[-1] & 0x0F
            code_int = struct.unpack(">I", h[offset_byte:offset_byte + 4])[0] & 0x7FFFFFFF
            expected = str(code_int % 1000000).zfill(6)
            if expected == code.strip():
                return True
        return False

    async def _record_attempt(self, username: str, ip_address: str, success: bool):
        """Record login attempt for brute-force tracking"""
        attempt = LoginAttempt(username=username, ip_address=ip_address, success=success)
        await self.attempt_repo.create(attempt)

    async def _audit(self, action: str, target_type: str, target_id: int, detail: str):
        log = AuditLog(action=action, target_type=target_type, target_id=target_id, detail=detail, ip_address="")
        await self.audit_repo.create(log)