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

import ipaddress

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

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


def _ip_in_allowlist(client_ip: str, allowed: list[str]) -> bool:
    """Check if client_ip matches any entry in the allowlist.

    Each entry can be:
      - Exact IP string: "1.2.3.4"
      - CIDR range: "10.0.0.0/8"
      - Hostname/domain: "proxy.example.com" (exact match only)
    """
    try:
        client_addr = ipaddress.ip_address(client_ip)
        for entry in allowed:
            entry = entry.strip()
            try:
                if "/" in entry:
                    if client_addr in ipaddress.ip_network(entry, strict=False):
                        return True
                else:
                    if client_addr == ipaddress.ip_address(entry):
                        return True
            except ValueError:
                # Not an IP — treat as hostname exact match
                if client_ip == entry:
                    return True
    except ValueError:
        # client_ip is not a valid IP (e.g., domain from reverse proxy)
        return client_ip in allowed
    return False

# 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 * 24 * 60  # 30 days
JWT_REFRESH_TOKEN_EXPIRE_DAYS = 30                # 30 days
JWT_WEBSSH_TOKEN_EXPIRE_MINUTES = 15

# Redis key prefix for multi-device refresh token storage
REFRESH_TOKEN_REDIS_PREFIX = "refresh_tokens:"


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": "..."}
        """
        # ── IP allowlist check ──
        # Only enforced when login_allowlist_enabled = "true"
        allowlist_on = (getattr(settings, "LOGIN_ALLOWLIST_ENABLED", "false") or "false").lower()
        if allowlist_on in ("true", "1", "yes", "on") and ip_address and ip_address not in ("", "unknown"):
            # Combine subscription IPs + manual IPs
            sub_ips_raw = (getattr(settings, "LOGIN_SUBSCRIPTION_IPS", "") or "").strip()
            manual_ips_raw = (getattr(settings, "LOGIN_MANUAL_IPS", "") or "").strip()
            all_ips_raw = ",".join(filter(None, [sub_ips_raw, manual_ips_raw]))
            allowed_ips = [ip.strip() for ip in all_ips_raw.split(",") if ip.strip()]
            if allowed_ips and not _ip_in_allowlist(ip_address, allowed_ips):
                await self._audit(
                    "login_ip_blocked", "admin", 0,
                    f"IP不在白名单: {ip_address} (用户: {username})",
                    ip_address=ip_address, admin_username=username,
                )
                return {
                    "success": False,
                    "reason": "ip_blocked",
                    "message": f"当前 IP ({ip_address}) 不在登录白名单中，请检查代理或联系管理员",
                }

        # 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"账号已锁定: {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 — update admin first (so JWT captures latest updated_at/token_version)
        admin.last_login = datetime.datetime.now(timezone.utc)
        await self.admin_repo.update(admin)

        # Generate JWT tokens AFTER admin update (captures latest updated_at)
        access_token = self._create_access_token(admin)
        refresh_token = self._create_refresh_token(admin)

        # Store refresh token in Redis (supports multi-device login)
        await self._store_refresh_token(admin.id, refresh_token)

        await self._record_attempt(username, ip_address, True)
        await self._audit(
            "login_success", "admin", admin.id, f"登录: {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 + multi-device support)

        Security: Refresh token format is ``token:admin_id:token_version``.
        Token version mismatch means this token is stale (e.g., password was changed),
        but does NOT invalidate other devices — only this token is rejected.
        Refresh tokens are stored in Redis Set (one admin can have multiple active sessions).
        """
        # 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, "刷新令牌格式无效",
                    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,
                    "刷新令牌引用的管理员不存在",
                    ip_address=ip_address,
                )
                return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}

            # Check token version — mismatch means this token is stale
            # (password changed, TOTP disabled, etc.), NOT a reuse attack.
            # Simply reject this token without punishing other devices.
            if admin.token_version != token_version:
                await self._audit(
                    "refresh_version_mismatch", "admin", admin.id,
                    f"刷新令牌版本不匹配: 令牌版本={token_version}, 当前版本={admin.token_version}",
                    ip_address=ip_address, admin_username=admin.username,
                )
                return {"success": False, "reason": "invalid_token", "message": "令牌已失效，请重新登录"}

            # Verify token exists in Redis (multi-device: Set membership check)
            if not await self._check_refresh_token(admin.id, refresh_token):
                await self._audit(
                    "refresh_token_not_found", "admin", admin.id,
                    "刷新令牌不在活跃会话中",
                    ip_address=ip_address, admin_username=admin.username,
                )
                return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}

        else:
            # Legacy format (old opaque token) — reject
            await self._audit(
                "refresh_legacy_format", "admin", 0, "旧版刷新令牌格式已废弃",
                ip_address=ip_address,
            )
            return {"success": False, "reason": "invalid_token", "message": "无效的刷新令牌"}

        # Rotate: remove old token from Redis, generate new token pair
        await self._remove_refresh_token(admin.id, refresh_token)

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

        # Store new refresh token in Redis
        await self._store_refresh_token(admin.id, new_refresh)

        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 ALL refresh tokens for an admin (e.g., 'logout everywhere').
        Increments token_version to also invalidate any cached access tokens.
        """
        admin = await self.admin_repo.get_by_id(admin_id)
        if admin:
            admin.token_version += 1  # Invalidate all existing tokens
            await self.admin_repo.update(admin)
            await self._delete_all_refresh_tokens(admin.id)
        return {"success": True, "message": "已登出所有设备"}

    async def logout_by_token(self, refresh_token: str) -> dict:
        """Invalidate a single refresh token (single-device logout).
        Other devices remain logged in.
        """
        # Try to parse new format token
        parts = refresh_token.rsplit(":", 2)
        admin = None
        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

            if admin:
                # Remove only this device's token from Redis
                await self._remove_refresh_token(admin.id, refresh_token)
        else:
            # Legacy format — try DB lookup, then clear all for safety
            admin = await self.admin_repo.get_by_refresh_token(refresh_token)
            if admin:
                await self._delete_all_refresh_tokens(admin.id)

        if admin:
            await self._audit(
                "logout", "admin", admin.id, "登出",
                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 设置已发起",
            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 已启用",
            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)
        # Invalidate all refresh tokens (all devices must re-login)
        await self._delete_all_refresh_tokens(admin.id)
        await self._audit(
            "disable_totp", "admin", admin_id, "TOTP 已禁用",
            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"为服务器 {server_id} 签发 WebSSH 令牌",
            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 'tv' (token_version) claim so tokens are immediately
        invalidated when password changes, TOTP is disabled, or token
        reuse is detected — no grace window.
        """
        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,
            "tv": admin.token_version or 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)

    # ── Redis Refresh Token Storage (multi-device) ──

    @staticmethod
    def _hash_token(token: str) -> str:
        """SHA-256 hash of refresh token for Redis storage (prevents token exposure if Redis is compromised)"""
        return hashlib.sha256(token.encode()).hexdigest()

    async def _store_refresh_token(self, admin_id: int, token: str):
        """Add a refresh token to Redis Set for this admin (supports multiple devices)"""
        try:
            from server.infrastructure.redis.client import get_redis
            redis = get_redis()
            key = f"{REFRESH_TOKEN_REDIS_PREFIX}{admin_id}"
            token_hash = self._hash_token(token)
            await redis.sadd(key, token_hash)
            # TTL = refresh token lifetime + small buffer
            await redis.expire(key, JWT_REFRESH_TOKEN_EXPIRE_DAYS * 86400 + 3600)
        except Exception as e:
            logger.error(f"Failed to store refresh token in Redis for admin {admin_id}: {e}")

    async def _check_refresh_token(self, admin_id: int, token: str) -> bool:
        """Check if a refresh token exists in Redis Set for this admin"""
        try:
            from server.infrastructure.redis.client import get_redis
            redis = get_redis()
            key = f"{REFRESH_TOKEN_REDIS_PREFIX}{admin_id}"
            token_hash = self._hash_token(token)
            return bool(await redis.sismember(key, token_hash))
        except Exception as e:
            logger.error(f"Failed to check refresh token in Redis for admin {admin_id}: {e}")
            return False

    async def _remove_refresh_token(self, admin_id: int, token: str):
        """Remove a single refresh token from Redis Set (single-device logout or rotation)"""
        try:
            from server.infrastructure.redis.client import get_redis
            redis = get_redis()
            key = f"{REFRESH_TOKEN_REDIS_PREFIX}{admin_id}"
            token_hash = self._hash_token(token)
            await redis.srem(key, token_hash)
        except Exception as e:
            logger.error(f"Failed to remove refresh token from Redis for admin {admin_id}: {e}")

    async def _delete_all_refresh_tokens(self, admin_id: int):
        """Delete all refresh tokens for an admin (logout everywhere / password change)"""
        try:
            from server.infrastructure.redis.client import get_redis
            redis = get_redis()
            key = f"{REFRESH_TOKEN_REDIS_PREFIX}{admin_id}"
            await redis.delete(key)
        except Exception as e:
            logger.error(f"Failed to delete all refresh tokens in Redis for admin {admin_id}: {e}")