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
import bcrypt
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
JWT_WEBSSH_TOKEN_EXPIRE_MINUTES = 15


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}",
                ip_address=ip_address, admin_username=username,
            )
            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(admin)

        # 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}",
            ip_address=ip_address, admin_username=admin.username,
        )

        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, ip_address: str = "") -> dict:
        """Exchange refresh token for new access token (with rotation + version-based reuse detection)

        Security: Refresh token format is ``token:admin_id:token_version``.
        If the version in the token doesn't match the admin's current version,
        it means this token was stolen/reused — invalidate all sessions for this admin.
        """
        # Parse new format token: "token:admin_id:token_version"
        parts = refresh_token.rsplit(":", 2)
        if len(parts) == 3:
            token_val, admin_id_str, token_version_str = parts
            try:
                admin_id = int(admin_id_str)
                token_version = int(token_version_str)
            except (ValueError, TypeError):
                await self._audit(
                    "refresh_invalid_format", "admin", 0, "Refresh token has invalid format",
                    ip_address=ip_address,
                )
                return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}

            # Get admin by ID (not by token)
            admin = await self.admin_repo.get_by_id(admin_id)
            if not admin:
                await self._audit(
                    "refresh_admin_not_found", "admin", admin_id,
                    "Refresh token references non-existent admin",
                    ip_address=ip_address,
                )
                return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}

            # Check token version — mismatch = reuse attack
            if admin.token_version != token_version:
                # Reuse attack detected! Invalidate all sessions by incrementing version
                original_version = admin.token_version
                admin.token_version += 1
                admin.jwt_refresh_token = None
                admin.jwt_token_expires = None
                await self.admin_repo.update(admin)
                await self._audit(
                    "refresh_reuse_attack", "admin", admin.id,
                    f"Token reuse detected! DB version={original_version}, token version={token_version}. "
                    f"All sessions invalidated (new version={admin.token_version}).",
                    ip_address=ip_address, admin_username=admin.username,
                )
                return {"success": False, "reason": "token_reuse", "message": "检测到令牌重用，所有会话已失效，请重新登录"}

            # Check if token matches current stored token (exact match required)
            if admin.jwt_refresh_token != refresh_token:
                await self._audit(
                    "refresh_token_mismatch", "admin", admin.id,
                    "Refresh token doesn't match stored token",
                    ip_address=ip_address, admin_username=admin.username,
                )
                return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}

        else:
            # Legacy format (old opaque token) — try exact match for backward compatibility
            admin = await self.admin_repo.get_by_refresh_token(refresh_token)
            if not admin:
                await self._audit(
                    "refresh_legacy_not_found", "admin", 0, "Legacy refresh token not found",
                    ip_address=ip_address,
                )
                return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}
            # Force rotation to new format on next refresh

        if admin.jwt_token_expires:
            expires = admin.jwt_token_expires
            now_cmp = datetime.datetime.now(timezone.utc)
            # MySQL returns tz-naive datetime; strip tzinfo before comparison to avoid TypeError
            if expires.tzinfo is None:
                now_cmp = now_cmp.replace(tzinfo=None)
            if expires < now_cmp:
                return {"success": False, "reason": "token_expired", "message": "刷新令牌已过期"}

        # Rotate: generate new token pair (old refresh token is invalidated by DB update)
        access_token = self._create_access_token(admin)
        new_refresh = self._create_refresh_token(admin)

        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 (also increments token_version to invalidate all sessions)"""
        admin = await self.admin_repo.get_by_id(admin_id)
        if admin:
            admin.token_version += 1  # Invalidate all existing refresh tokens
            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)"""
        # Try to parse new format token
        parts = refresh_token.rsplit(":", 2)
        if len(parts) == 3:
            admin_id_str = parts[1]
            try:
                admin_id = int(admin_id_str)
                admin = await self.admin_repo.get_by_id(admin_id)
            except (ValueError, TypeError):
                admin = None
        else:
            # Legacy format
            admin = await self.admin_repo.get_by_refresh_token(refresh_token)

        if admin:
            admin.token_version += 1  # Invalidate all existing refresh tokens
            admin.jwt_refresh_token = None
            admin.jwt_token_expires = None
            await self.admin_repo.update(admin)
            await self._audit(
                "logout", "admin", admin.id, "Logout",
                admin_username=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
        from urllib.parse import quote
        issuer = settings.SYSTEM_NAME
        issuer_enc = quote(issuer, safe="")
        user_enc = quote(admin.username, safe="")
        uri = f"otpauth://totp/{issuer_enc}:{user_enc}?secret={secret}&issuer={issuer_enc}"

        await self._audit(
            "setup_totp", "admin", admin_id, "TOTP setup initiated",
            admin_username=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, "TOTP enabled",
            admin_username=admin.username,
        )
        return {"success": True, "message": "TOTP已启用"}

    async def disable_totp(
        self,
        admin_id: int,
        current_password: str,
        totp_code: Optional[str] = None,
    ) -> dict:
        """Disable TOTP — requires current password and TOTP code when enabled."""
        admin = await self.admin_repo.get_by_id(admin_id)
        if not admin:
            return {"success": False, "reason": "admin_not_found"}

        if not self._verify_password(current_password, admin.password_hash):
            return {"success": False, "reason": "invalid_password", "message": "当前密码错误"}

        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):
                return {"success": False, "reason": "invalid_totp", "message": "TOTP验证码错误"}

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

    async def create_webssh_token(self, admin: Admin, server_id: int) -> dict:
        """Issue a short-lived JWT scoped to one server for WebSSH.

        General access tokens must not be used for terminal WebSocket (no server binding).
        """
        token = self._create_webssh_token(admin, server_id)
        await self._audit(
            "webssh_token",
            "server",
            server_id,
            f"WebSSH token issued for server {server_id}",
            admin_username=admin.username,
        )
        return {
            "success": True,
            "webssh_token": token,
            "server_id": server_id,
            "expires_in": JWT_WEBSSH_TOKEN_EXPIRE_MINUTES * 60,
        }

    # ── JWT Token Helpers ──

    def _create_access_token(self, admin: Admin) -> str:
        """Create JWT access token with admin claims

        Includes 'updated_at' claim so tokens are invalidated when admin changes password.
        """
        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),
            "updated": int(admin.updated_at.timestamp()) if admin.updated_at else 0,
        }
        return jwt.encode(payload, settings.SECRET_KEY, algorithm="HS256")

    def _create_webssh_token(self, admin: Admin, server_id: int) -> str:
        """JWT for WebSSH only — binds admin to a single server_id."""
        import jwt
        from server.config import settings

        now = datetime.datetime.now(timezone.utc)
        payload = {
            "sub": str(admin.id),
            "username": admin.username,
            "server_id": int(server_id),
            "purpose": "webssh",
            "iat": now,
            "exp": now + datetime.timedelta(minutes=JWT_WEBSSH_TOKEN_EXPIRE_MINUTES),
            "updated": int(admin.updated_at.timestamp()) if admin.updated_at else 0,
        }
        return jwt.encode(payload, settings.SECRET_KEY, algorithm="HS256")

    def _create_refresh_token(self, admin: Admin) -> str:
        """Create refresh token with version binding: token:admin_id:token_version

        Format includes admin_id and token_version so we can detect reuse attacks
        even if the token is not found in DB (e.g., after rotation).
        """
        token = secrets.token_urlsafe(32)
        return f"{token}:{admin.id}:{admin.token_version}"

    def _decode_access_token(self, token: str) -> Optional[dict]:
        """Decode and verify JWT access token

        Validates that required claims (exp, sub) are present after decoding.
        PyJWT checks exp automatically, but we also ensure the claim exists
        to reject tokens with missing or null expiry.
        """
        try:
            import jwt
            from server.config import settings
            payload = jwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"])
            # Require both 'exp' and 'sub' claims — reject malformed tokens
            if "exp" not in payload or "sub" not in payload:
                return None
            return payload
        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 bcrypt hash"""
        return bcrypt.checkpw(password.encode(), password_hash.encode())

    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,
        ip_address: str = "",
        admin_username: str = "",
    ):
        log = AuditLog(
            action=action,
            target_type=target_type,
            target_id=target_id,
            detail=detail,
            ip_address=ip_address or None,
            admin_username=admin_username or None,
        )
        await self.audit_repo.create(log)