zsglpt/security/blacklist.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations

import threading
from datetime import timedelta
from typing import List, Optional

import db_pool
from db.utils import get_cst_now, get_cst_now_str


class BlacklistManager:
    """黑名单管理器"""

    def __init__(self) -> None:
        self._schema_ready = False
        self._schema_lock = threading.Lock()

    def is_ip_banned(self, ip: str) -> bool:
        """检查IP是否被封禁"""
        ip_text = str(ip or "").strip()[:64]
        if not ip_text:
            return False
        now_str = get_cst_now_str()

        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT 1
                FROM ip_blacklist
                WHERE ip = ?
                  AND is_active = 1
                  AND (expires_at IS NULL OR expires_at > ?)
                LIMIT 1
                """,
                (ip_text, now_str),
            )
            return cursor.fetchone() is not None

    def is_user_banned(self, user_id: int) -> bool:
        """检查用户是否被封禁"""
        if user_id is None:
            return False
        self._ensure_schema()
        user_id_int = int(user_id)
        now_str = get_cst_now_str()

        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT 1
                FROM user_blacklist
                WHERE user_id = ?
                  AND is_active = 1
                  AND (expires_at IS NULL OR expires_at > ?)
                LIMIT 1
                """,
                (user_id_int, now_str),
            )
            return cursor.fetchone() is not None

    def ban_ip(self, ip: str, reason: str, duration_hours: int = 24, permanent: bool = False):
        """封禁IP"""
        ip_text = str(ip or "").strip()[:64]
        if not ip_text:
            return False
        reason_text = str(reason or "").strip()[:512]
        now_str = get_cst_now_str()

        expires_at: Optional[str]
        if permanent:
            expires_at = None
        else:
            hours = max(1, int(duration_hours))
            expires_at = (get_cst_now() + timedelta(hours=hours)).strftime("%Y-%m-%d %H:%M:%S")

        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                INSERT INTO ip_blacklist (ip, reason, is_active, added_at, expires_at)
                VALUES (?, ?, 1, ?, ?)
                ON CONFLICT(ip) DO UPDATE SET
                    reason = excluded.reason,
                    is_active = 1,
                    added_at = excluded.added_at,
                    expires_at = excluded.expires_at
                """,
                (ip_text, reason_text, now_str, expires_at),
            )
            conn.commit()
            return True

    def ban_user(self, user_id: int, reason: str):
        """封禁用户"""
        return self._ban_user_internal(user_id, reason=reason, duration_hours=24, permanent=False)

    def unban_ip(self, ip: str):
        """解除IP封禁"""
        ip_text = str(ip or "").strip()[:64]
        if not ip_text:
            return False
        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute("UPDATE ip_blacklist SET is_active = 0 WHERE ip = ?", (ip_text,))
            conn.commit()
            return cursor.rowcount > 0

    def unban_user(self, user_id: int):
        """解除用户封禁"""
        if user_id is None:
            return False
        self._ensure_schema()
        user_id_int = int(user_id)
        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute("UPDATE user_blacklist SET is_active = 0 WHERE user_id = ?", (user_id_int,))
            conn.commit()
            return cursor.rowcount > 0

    def get_banned_ips(self) -> List[dict]:
        """获取所有被封禁的IP"""
        now_str = get_cst_now_str()
        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT ip, reason, is_active, added_at, expires_at
                FROM ip_blacklist
                WHERE is_active = 1
                  AND (expires_at IS NULL OR expires_at > ?)
                ORDER BY added_at DESC
                """,
                (now_str,),
            )
            return [dict(row) for row in cursor.fetchall()]

    def get_banned_users(self) -> List[dict]:
        """获取所有被封禁的用户"""
        self._ensure_schema()
        now_str = get_cst_now_str()
        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT user_id, reason, is_active, added_at, expires_at
                FROM user_blacklist
                WHERE is_active = 1
                  AND (expires_at IS NULL OR expires_at > ?)
                ORDER BY added_at DESC
                """,
                (now_str,),
            )
            return [dict(row) for row in cursor.fetchall()]

    def cleanup_expired(self):
        """清理过期的封禁记录"""
        now_str = get_cst_now_str()
        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                UPDATE ip_blacklist
                SET is_active = 0
                WHERE is_active = 1
                  AND expires_at IS NOT NULL
                  AND expires_at <= ?
                """,
                (now_str,),
            )
            conn.commit()

        self._ensure_schema()
        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                UPDATE user_blacklist
                SET is_active = 0
                WHERE is_active = 1
                  AND expires_at IS NOT NULL
                  AND expires_at <= ?
                """,
                (now_str,),
            )
            conn.commit()

    # ==================== Internal ====================

    def _ensure_schema(self) -> None:
        if self._schema_ready:
            return
        with self._schema_lock:
            if self._schema_ready:
                return
            with db_pool.get_db() as conn:
                cursor = conn.cursor()
                cursor.execute(
                    """
                    CREATE TABLE IF NOT EXISTS user_blacklist (
                        user_id INTEGER PRIMARY KEY,
                        reason TEXT,
                        is_active INTEGER DEFAULT 1,
                        added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        expires_at TIMESTAMP
                    )
                    """
                )
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_blacklist_active ON user_blacklist(is_active)")
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_blacklist_expires ON user_blacklist(expires_at)")
                conn.commit()
            self._schema_ready = True

    def _ban_user_internal(
        self,
        user_id: int,
        *,
        reason: str,
        duration_hours: int = 24,
        permanent: bool = False,
    ) -> bool:
        if user_id is None:
            return False
        self._ensure_schema()
        user_id_int = int(user_id)
        reason_text = str(reason or "").strip()[:512]
        now_str = get_cst_now_str()

        expires_at: Optional[str]
        if permanent:
            expires_at = None
        else:
            hours = max(1, int(duration_hours))
            expires_at = (get_cst_now() + timedelta(hours=hours)).strftime("%Y-%m-%d %H:%M:%S")

        with db_pool.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                INSERT INTO user_blacklist (user_id, reason, is_active, added_at, expires_at)
                VALUES (?, ?, 1, ?, ?)
                ON CONFLICT(user_id) DO UPDATE SET
                    reason = excluded.reason,
                    is_active = 1,
                    added_at = excluded.added_at,
                    expires_at = excluded.expires_at
                """,
                (user_id_int, reason_text, now_str, expires_at),
            )
            conn.commit()
            return True