SenSu/services/web_panel/routes/auth.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import secrets
import time
import logging
from aiohttp import web
from ..utils.auth import panel_auth

logger = logging.getLogger(__name__)

# 全局 Session 存储 (内存型)
PANEL_SESSION_STORE = {}

# 登录频率限制 — {ip: [fail_count, lock_until_timestamp]}
_LOGIN_FAILS: dict[str, list] = {}
_MAX_FAILS = 5
_LOCK_SECONDS = 60


def _check_rate_limit(ip: str) -> bool:
    """检查 IP 是否被限流。返回 True = 允许尝试"""
    now = time.time()
    entry = _LOGIN_FAILS.get(ip)
    if entry:
        fail_count, lock_until = entry
        if lock_until and now < lock_until:
            return False  # still locked
        if lock_until and now >= lock_until:
            _LOGIN_FAILS.pop(ip, None)  # lock expired, reset
    return True


def _record_fail(ip: str):
    now = time.time()
    entry = _LOGIN_FAILS.get(ip)
    if entry is None:
        entry = [0, None]
    entry[0] += 1
    if entry[0] >= _MAX_FAILS:
        entry[1] = now + _LOCK_SECONDS
        logger.warning(f"🔒 IP {ip} 登录锁定 {_LOCK_SECONDS}s ({_MAX_FAILS} 次失败)")
    else:
        logger.debug(f"IP {ip} 登录失败计数: {entry[0]}/{_MAX_FAILS}")
    _LOGIN_FAILS[ip] = entry


def _clear_fails(ip: str):
    _LOGIN_FAILS.pop(ip, None)


def _load_sessions_from_db(app):
    """从数据库恢复持久化的 session"""
    try:
        db = app.get("sensu_db")
        if not db:
            sm = app.get("service_manager")
            if sm:
                try:
                    db = sm.get_service("sensu_db")
                except Exception:
                    pass
        if db:
            raw = db.get_config("panel_sessions", "{}")
            stored = json.loads(raw) if raw else {}
            for token, info in stored.items():
                PANEL_SESSION_STORE[token] = info
            if stored:
                logger.info(f"📦 从数据库恢复了 {len(stored)} 个会话")
    except Exception as e:
        logger.warning(f"会话恢复失败: {e}")


def _save_sessions_to_db(app):
    """持久化当前 session 到数据库"""
    try:
        db = app.get("sensu_db")
        if not db:
            sm = app.get("service_manager")
            if sm:
                try:
                    db = sm.get_service("sensu_db")
                except Exception:
                    pass
        if db:
            db.set_config("panel_sessions", json.dumps(PANEL_SESSION_STORE))
    except Exception as e:
        logger.debug(f"会话持久化失败: {e}")


def setup_routes(app, prefix=''):
    """注册面板认证路由"""
    import json as _json

    # 🟢 关键：将 Session Store 挂载到 app，供拦截器读取
    app['panel_session_store'] = PANEL_SESSION_STORE

    # 从数据库恢复持久化会话
    _load_sessions_from_db(app)

    # 路由注册
    app.router.add_post(f'{prefix}/api/login', handle_login)
    app.router.add_post(f'{prefix}/api/logout', panel_auth(handle_logout))
    app.router.add_get(f'{prefix}/api/auth/status', panel_auth(handle_auth_status))

async def handle_login(req):
    """处理面板登录"""
    try:
        # 频率限制
        ip = req.remote
        if not _check_rate_limit(ip):
            return web.json_response(
                {"success": False, "msg": "尝试次数过多，请 60 秒后重试"},
                status=429,
            )

        data = await req.json()
        username = data.get('username')
        password = data.get('password')

        cfg = req.app.get('panel_config', {})
        cfg_user = cfg.get('username', 'admin')
        cfg_pass = cfg.get('password', 'admin')

        # 校验配置中的账号密码
        if username == cfg_user and password == cfg_pass:
            _clear_fails(ip)
            # 登录成功：生成 Token
            token = secrets.token_hex(16)
            
            # 写入 Session Store
            user_info = {
                "username": username,
                "perms": ["admin"],
                "login_time": __import__('time').time()
            }
            PANEL_SESSION_STORE[token] = user_info
            _save_sessions_to_db(req.app)  # 持久化到数据库

            logger.info(f"✅ 面板登录成功: {username} (Session: {token[:4]}...)")
            
            resp = web.json_response({"success": True, "username": username})
            # 设置 Cookie
            resp.set_cookie("panel_token", token, max_age=259200, httponly=True, samesite="Lax")
            return resp
        else:
            _record_fail(ip)
            logger.warning(f"❌ 面板登录失败: 用户 {username} 密码错误 (IP: {ip})")
            return web.json_response({"success": False, "msg": "用户名或密码错误"}, status=401)
            
    except Exception as e:
        logger.error(f"登录异常: {e}")
        return web.json_response({"error": str(e)}, status=500)

async def handle_logout(req):
    """处理退出登录"""
    token = req.cookies.get("panel_token")
    if token and token in PANEL_SESSION_STORE:
        del PANEL_SESSION_STORE[token]
        _save_sessions_to_db(req.app)  # 持久化删除
        logger.info(f"👋 用户退出登录")

    resp = web.json_response({"success": True})
    resp.del_cookie("panel_token")
    return resp

async def handle_auth_status(req):
    """获取当前认证状态 (被 panel_auth 拦截，能进来说明已认证)"""
    user = req.get('user', {})
    return web.json_response({
        "authenticated": True,
        "username": user.get("username", "Unknown"),
        "perms": user.get("perms", [])
    })