diff --git a/.env.example b/.env.example index 777e235..6277596 100644 --- a/.env.example +++ b/.env.example @@ -34,6 +34,12 @@ DEFAULT_NEGATIVE=low quality, blurry, bad anatomy, extra digits # Доступ: пусто — всем; иначе список id через запятую ALLOWED_USER_IDS= +# Админы: их промпты не фильтруются +ADMIN_USER_IDS= +MODERATION_ENABLED=true +# Дополнительные имена публичных людей через запятую +BLOCKED_PUBLIC_FIGURES= + # Веб-история: только локальная сеть, без авторизации WEB_BIND_IP=192.168.31.158 WEB_PAGE_SIZE=50 diff --git a/README.md b/README.md index 9409514..7c75b30 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ Telegram-бот для генерации и редактирования изо - **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки. - **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений. - **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI. +- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - SQLite: запросы, статусы, seed, пути к файлам. @@ -94,6 +95,32 @@ Workflow: `.gitea/workflows/docker-build.yml`. Runner использует label Отдельную модель скачивать не нужно, используется существующий CLIP-файл ComfyUI. +## Фильтрация промптов + +Фильтр работает для всех, кроме пользователей из `ADMIN_USER_IDS`: + +```dotenv +MODERATION_ENABLED=true +ADMIN_USER_IDS=7467236616 +BLOCKED_PUBLIC_FIGURES= +``` + +`BLOCKED_PUBLIC_FIGURES` добавляет имена к встроенному списку (через запятую). + +Что блокируется: + +- CSAM: child/kid/minor/teen/young/loli/shota/school/underage + NSFW; +- зоофилия: sex/nude + animal/dog/cat/horse/beast; +- публичные люди + nude/naked/sex/undress/deepfake; +- undress/nude/naked + photo/real/person/woman/man; +- насилие + NSFW; +- экстремизм + инструкции/символика; +- обходные написания (`n_u_d_e`, `l0li`, `ch1ld` и т.п.). + +Для i2i OpenCV проверяет наличие лица/человека на входном фото: NSFW-промпт в этом случае блокируется. + +При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`. + ## Веб-история После `docker compose up -d` доступна страница: diff --git a/bot.py b/bot.py index 63781e1..a5d2023 100644 --- a/bot.py +++ b/bot.py @@ -20,6 +20,7 @@ from telegram.ext import ( from comfy_client import ComfyClient from config import Settings, setup_logging from db import Database, Request, UserAlreadyQueued +from moderation import PromptBlocked, detect_person from prompt_enhancer import PromptEnhancer from queue_manager import QueueEvent, QueueFull, QueueManager from workflows import WorkflowBuilder @@ -32,6 +33,7 @@ START_TEXT = ( "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n" "/status — узнать состояние текущей задачи." ) +BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны." class MpakoBot: @@ -294,6 +296,16 @@ class MpakoBot: if message is None or user is None or self._queue is None: return False + is_admin = user.id in self._settings.admin_user_ids + has_person = False + if ( + mode == "i2i" + and input_path + and not is_admin + and self._settings.moderation_enabled + ): + has_person = await asyncio.to_thread(detect_person, input_path) + try: request, position = await self._queue.enqueue( user_id=user.id, @@ -301,7 +313,12 @@ class MpakoBot: mode=mode, prompt=prompt, input_file_path=input_path, + has_person=has_person, + is_admin=is_admin, ) + except PromptBlocked: + await message.reply_text(BLOCKED_TEXT) + return False except UserAlreadyQueued: active = await self._db.get_active_by_user(user.id) text = "Ты уже в очереди." diff --git a/checklist.md b/checklist.md index da0659c..3be93ac 100644 --- a/checklist.md +++ b/checklist.md @@ -6,10 +6,10 @@ ## Текущий статус -- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13) и улучшение промптов (14). -- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, переключение T2I → I2I. -- 25 тестов проходят: БД, workflows, очередь, веб-история, enhancer, вспомогательные функции ComfyUI. -- Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена). +- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14) и модерация (15). +- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. +- 45 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, вспомогательные функции ComfyUI. +- Не проверено: очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. ## 0. Допущения и решения по умолчанию @@ -318,6 +318,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); **DoD:** русский запрос превращается в английский промпт, изображение генерируется по нему; при сбое LLM используется оригинал. +## 15. Фильтрация промптов + +- [x] `is_prompt_allowed`: нормализация, hard/context-блоки, обходные паттерны. +- [x] Админы из `ADMIN_USER_IDS` обходят фильтр. +- [x] Логирование блокировок в `blocked_requests`. +- [x] Ответ «Запрос нарушает правила. Звёзды не списаны.». +- [x] i2i: детекция лица/человека через OpenCV. +- [x] Тесты правил и очереди. + +**DoD:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра. + ## Порядок работ 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. diff --git a/config.py b/config.py index 89a3c64..56eca36 100644 --- a/config.py +++ b/config.py @@ -59,6 +59,16 @@ def _user_ids(name: str) -> frozenset[int]: return frozenset(values) +def _csv(name: str) -> frozenset[str]: + raw = os.getenv(name, "") + values = { + part.strip().lower() + for part in raw.replace(";", ",").split(",") + if part.strip() + } + return frozenset(values) + + @dataclass(frozen=True) class Settings: bot_token: str @@ -75,12 +85,15 @@ class Settings: worker_concurrency: int = 1 max_queue_size: int = 20 allowed_user_ids: frozenset[int] = field(default_factory=frozenset) + admin_user_ids: frozenset[int] = field(default_factory=frozenset) progress_interval: float = 3.0 generation_timeout: float = 600.0 log_level: str = "INFO" default_negative: str = "low quality, blurry, bad anatomy, extra digits" max_prompt_length: int = 1000 max_photo_bytes: int = 15 * 1024 * 1024 + moderation_enabled: bool = True + blocked_public_figures: frozenset[str] = field(default_factory=frozenset) web_page_size: int = 50 web_bind_ip: str = "127.0.0.1" @@ -129,6 +142,7 @@ class Settings: worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)), max_queue_size=max(1, _int("MAX_QUEUE_SIZE", 20)), allowed_user_ids=_user_ids("ALLOWED_USER_IDS"), + admin_user_ids=_user_ids("ADMIN_USER_IDS"), progress_interval=max(1.0, _float("PROGRESS_INTERVAL", 3.0)), generation_timeout=max(30.0, _float("GENERATION_TIMEOUT", 600.0)), log_level=os.getenv("LOG_LEVEL", "INFO").strip().upper() or "INFO", @@ -136,6 +150,8 @@ class Settings: or cls.default_negative, max_prompt_length=max(1, _int("MAX_PROMPT_LENGTH", 1000)), max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)), + moderation_enabled=_bool("MODERATION_ENABLED", True), + blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"), web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))), web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip() or "127.0.0.1", diff --git a/db.py b/db.py index b45eac0..21bcb6e 100644 --- a/db.py +++ b/db.py @@ -78,6 +78,29 @@ class UserStat: return f"@{self.username}" if self.username else str(self.user_id) +@dataclass(frozen=True) +class BlockedRequest: + id: int + user_id: int + timestamp: str + prompt: str + reason: str + mode: str | None + has_person: bool + + @classmethod + def from_row(cls, row: aiosqlite.Row) -> BlockedRequest: + return cls( + id=row["id"], + user_id=row["user_id"], + timestamp=row["timestamp"], + prompt=row["prompt"], + reason=row["reason"], + mode=row["mode"], + has_person=bool(row["has_person"]), + ) + + SCHEMA = """ CREATE TABLE IF NOT EXISTS requests ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -105,6 +128,19 @@ CREATE INDEX IF NOT EXISTS idx_requests_status_id CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active ON requests(user_id) WHERE status IN ('queued', 'running'); + +CREATE TABLE IF NOT EXISTS blocked_requests ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER NOT NULL, + timestamp TEXT NOT NULL, + prompt TEXT NOT NULL, + reason TEXT NOT NULL, + mode TEXT, + has_person INTEGER NOT NULL DEFAULT 0 +); + +CREATE INDEX IF NOT EXISTS idx_blocked_user_time + ON blocked_requests(user_id, id); """ @@ -378,6 +414,36 @@ class Database: (error_text[:1000], utc_now(), request_id), ) + async def log_blocked( + self, + *, + user_id: int, + prompt: str, + reason: str, + mode: str | None = None, + has_person: bool = False, + ) -> int: + conn = self._connection() + async with self._write_lock: + cursor = await conn.execute( + """ + INSERT INTO blocked_requests + (user_id, timestamp, prompt, reason, mode, has_person) + VALUES (?, ?, ?, ?, ?, ?) + """, + (user_id, utc_now(), prompt, reason, mode, int(has_person)), + ) + await conn.commit() + return int(cursor.lastrowid) + + async def list_blocked(self, limit: int = 100) -> list[BlockedRequest]: + conn = self._connection() + async with conn.execute( + "SELECT * FROM blocked_requests ORDER BY id DESC LIMIT ?", (limit,) + ) as cursor: + rows = await cursor.fetchall() + return [BlockedRequest.from_row(row) for row in rows] + async def recover(self, reason: str = "прервано перезапуском") -> list[Request]: """Помечает running как failed и возвращает queued для повторной постановки.""" conn = self._connection() diff --git a/moderation.py b/moderation.py new file mode 100644 index 0000000..b2733d4 --- /dev/null +++ b/moderation.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +import logging +import re +from pathlib import Path + +logger = logging.getLogger(__name__) + +DIGIT_MAP = str.maketrans( + { + "0": "o", + "1": "i", + "3": "e", + "4": "a", + "5": "s", + "7": "t", + } +) + +CSAM_TERMS = { + "child", + "kid", + "minor", + "teen", + "young", + "loli", + "shota", + "school", + "underage", + "ребен", + "ребён", + "дети", + "детск", + "школь", + "малолет", + "подрост", +} +NSFW_TERMS = { + "nude", + "naked", + "sex", + "erotic", + "explicit", + "porn", + "nsfw", + "нюд", + "голый", + "голая", + "голое", + "голые", + "голую", + "голой", + "голых", + "секс", + "эрот", + "порн", + "обнаж", + "раздет", +} +ANIMAL_TERMS = { + "animal", + "dog", + "cat", + "horse", + "beast", + "животн", + "собак", + "кошк", + "лошад", + "звер", +} +PUBLIC_FIGURES = { + "trump", + "biden", + "putin", + "zelensky", + "obama", + "clinton", + "harris", + "macron", + "merkel", + "modi", + "netanyahu", + "kimjongun", + "elonmusk", + "taylorswift", + "beyonce", + "kanyewest", + "ronaldo", + "messi", + "madonna", + "bieber", + "kardashian", + "spears", + "michaeljackson", + "queenelizabeth", + "kingcharles", + "princewilliam", + "katemiddleton", + "трамп", + "байден", + "путин", + "зеленский", + "обама", + "клинтон", + "маск", +} +REAL_PERSON_TERMS = { + "nude", + "naked", + "sex", + "undress", + "deepfake", + "нюд", + "голый", + "голая", + "голое", + "голые", + "голую", + "голой", + "голых", + "секс", + "раздет", + "дипфейк", +} +NONCONSENSUAL_TERMS = { + "undress", + "nude", + "naked", + "раздет", + "нюд", + "голый", + "голая", + "голое", + "голые", + "голую", + "голой", + "голых", +} +PERSON_TERMS = { + "photo", + "real", + "person", + "woman", + "man", + "фото", + "реальн", + "человек", + "женщин", + "мужчин", + "девуш", + "парн", +} +VIOLENCE_TERMS = { + "gore", + "snuff", + "torture", + "dismember", + "blood", + "кров", + "пытк", + "расчлен", +} +EXTREMISM_TERMS = { + "terrorist", + "nazi", + "isis", + "jihad", + "bomb", + "террорист", + "нацист", + "игил", + "джихад", + "бомб", +} +INSTRUCTION_TERMS = { + "how", + "instruction", + "instructions", + "guide", + "tutorial", + "symbol", + "symbols", + "flag", + "uniform", + "manifesto", + "recruit", + "attack", + "make", + "build", + "инструкц", + "символ", + "флаг", + "как", + "сделать", + "создать", + "собрать", +} + +OBFUSCATED_PATTERN = re.compile( + r"l[0o1]l[i1]|sh[0o]t[a4]|ch[i1]ld|n[uv]de|undr[e3]ss", + re.IGNORECASE, +) + + +class PromptBlocked(RuntimeError): + def __init__(self, reason: str) -> None: + super().__init__(f"Промпт заблокирован: {reason}") + self.reason = reason + + +def _normalize(prompt: str) -> tuple[str, str, set[str]]: + lowered = prompt.lower() + normalized = lowered.translate(DIGIT_MAP) + compact = re.sub(r"[\s_\-]+", "", normalized) + + tokens = re.split(r"[\s_\-]+", normalized) + merged: list[str] = [] + buffer = "" + for token in tokens: + if not token: + continue + if len(token) == 1: + buffer += token + continue + if buffer: + merged.append(buffer) + buffer = "" + merged.append(token) + if buffer: + merged.append(buffer) + return lowered, compact, set(merged) + + +def _has(term: str, compact: str, words: set[str]) -> bool: + return term in compact or term in words + + +def is_prompt_allowed( + prompt: str, + *, + mode: str = "t2i", + has_person: bool = False, + is_admin: bool = False, + extra_public_figures: set[str] | frozenset[str] | None = None, +) -> tuple[bool, str | None]: + """Проверка промпта. Возвращает (allowed, reason_code).""" + if is_admin: + return True, None + + lowered, compact, words = _normalize(prompt) + figures = PUBLIC_FIGURES + if extra_public_figures: + extra = { + re.sub(r"[\s_\-]+", "", name.lower().translate(DIGIT_MAP)) + for name in extra_public_figures + if name.strip() + } + figures = PUBLIC_FIGURES | extra + nsfw = any(_has(term, compact, words) for term in NSFW_TERMS) + csam = any(_has(term, compact, words) for term in CSAM_TERMS) + + if csam and nsfw: + return False, "csam" + if _has("loli", compact, words) or _has("shota", compact, words): + return False, "csam" + + if ( + any(_has(term, compact, words) for term in ("sex", "nude", "naked")) + and any(_has(term, compact, words) for term in ANIMAL_TERMS) + ): + return False, "bestiality" + + if any(name in compact for name in figures) and any( + _has(term, compact, words) for term in REAL_PERSON_TERMS + ): + return False, "real_person" + + if any(_has(term, compact, words) for term in NONCONSENSUAL_TERMS) and any( + _has(term, compact, words) for term in PERSON_TERMS + ): + return False, "nonconsensual_sexualization" + + if mode == "i2i" and has_person and any( + _has(term, compact, words) for term in ("nude", "naked", "undress", "sex") + ): + return False, "i2i_person_nsfw" + + if any(_has(term, compact, words) for term in VIOLENCE_TERMS) and nsfw: + return False, "violence_sexual" + + if any(_has(term, compact, words) for term in EXTREMISM_TERMS) and any( + _has(term, compact, words) for term in INSTRUCTION_TERMS + ): + return False, "extremism" + + if re.search(r"[0-9_\-]", lowered) and OBFUSCATED_PATTERN.search(compact): + return False, "obfuscated_nsfw" + + return True, None + + +def detect_person(path: str | Path) -> bool: + """Ищет лицо или человека через OpenCV. При ошибке считает, что человека нет.""" + try: + import cv2 # type: ignore[import-not-found] + except ImportError: + logger.warning("OpenCV не установлен, детекция людей отключена") + return False + + try: + image = cv2.imread(str(path)) + if image is None: + return False + + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + cascade = cv2.CascadeClassifier( + str(Path(cv2.data.haarcascades) / "haarcascade_frontalface_default.xml") + ) + faces = cascade.detectMultiScale( + gray, + scaleFactor=1.1, + minNeighbors=5, + minSize=(24, 24), + ) + if len(faces) > 0: + return True + + height, width = image.shape[:2] + scale = min(1.0, 640.0 / max(height, width)) + if scale < 1.0: + image = cv2.resize(image, (int(width * scale), int(height * scale))) + hog = cv2.HOGDescriptor() + hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector()) + _, weights = hog.detectMultiScale( + image, + winStride=(8, 8), + padding=(8, 8), + scale=1.05, + ) + return any(float(weight) > 0.5 for weight in weights) + except Exception: + logger.warning("Ошибка детекции человека в %s", path, exc_info=True) + return False diff --git a/queue_manager.py b/queue_manager.py index 2c977b3..b515b90 100644 --- a/queue_manager.py +++ b/queue_manager.py @@ -13,6 +13,7 @@ from pathlib import Path from comfy_client import ComfyClient, ComfyError from config import Settings from db import Database, Request, UserAlreadyQueued +from moderation import PromptBlocked, is_prompt_allowed from prompt_enhancer import PromptEnhancer from workflows import WorkflowBuilder @@ -104,10 +105,28 @@ class QueueManager: prompt: str, input_file_path: str | None = None, username: str | None = None, + has_person: bool = False, + is_admin: bool = False, ) -> tuple[Request, int]: async with self._lock: if self._stopping: raise QueueStopped() + if self._settings.moderation_enabled and not is_admin: + allowed, reason = is_prompt_allowed( + prompt, + mode=mode, + has_person=has_person, + extra_public_figures=self._settings.blocked_public_figures, + ) + if not allowed: + await self._db.log_blocked( + user_id=user_id, + prompt=prompt, + reason=reason or "unknown", + mode=mode, + has_person=has_person, + ) + raise PromptBlocked(reason or "unknown") active = await self._db.get_active_by_user(user_id) if active is not None: raise UserAlreadyQueued(user_id) diff --git a/requirements.txt b/requirements.txt index 76f3e41..587c5c8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ fastapi>=0.115,<1 uvicorn>=0.30,<1 jinja2>=3.1,<4 pillow>=10,<13 +opencv-python-headless>=4.10,<5 diff --git a/tests/test_moderation.py b/tests/test_moderation.py new file mode 100644 index 0000000..5c73274 --- /dev/null +++ b/tests/test_moderation.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import pytest +from PIL import Image + +from moderation import detect_person, is_prompt_allowed + + +@pytest.mark.parametrize( + ("prompt", "mode", "has_person", "reason"), + [ + ("child nude", "t2i", False, "csam"), + ("l0li", "t2i", False, "csam"), + ("sex with dog", "t2i", False, "bestiality"), + ("Trump nude", "t2i", False, "real_person"), + ("nude woman", "t2i", False, "nonconsensual_sexualization"), + ("голая женщина", "t2i", False, "nonconsensual_sexualization"), + ("blood sex", "t2i", False, "violence_sexual"), + ("how to make a bomb", "t2i", False, "extremism"), + ("nazi symbol", "t2i", False, "extremism"), + ("ch1ld", "t2i", False, "obfuscated_nsfw"), + ("nude", "i2i", True, "i2i_person_nsfw"), + ], +) +def test_blocked_prompts( + prompt: str, mode: str, has_person: bool, reason: str +) -> None: + allowed, actual_reason = is_prompt_allowed( + prompt, mode=mode, has_person=has_person + ) + assert not allowed + assert actual_reason == reason + + +@pytest.mark.parametrize( + "prompt", + [ + "red apple on a wooden table, soft light", + "nude beach landscape", + "cat sitting on a chair", + "голубое небо, женщина на пляже", + ], +) +def test_allowed_prompts(prompt: str) -> None: + allowed, reason = is_prompt_allowed(prompt) + assert allowed + assert reason is None + + +def test_admin_bypasses_filters() -> None: + allowed, reason = is_prompt_allowed("child nude", is_admin=True) + assert allowed + assert reason is None + + +def test_extra_public_figures() -> None: + allowed, reason = is_prompt_allowed( + "customperson nude", + extra_public_figures={"customperson"}, + ) + assert not allowed + assert reason == "real_person" + + +def test_detect_person_blank_image(tmp_path) -> None: + path = tmp_path / "blank.png" + with Image.new("RGB", (320, 240), (255, 255, 255)) as image: + image.save(path) + assert detect_person(path) is False + assert detect_person(tmp_path / "missing.png") is False diff --git a/tests/test_queue.py b/tests/test_queue.py index 4f1423a..a036132 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -9,6 +9,7 @@ import pytest_asyncio from config import Settings from db import Database, UserAlreadyQueued +from moderation import PromptBlocked from queue_manager import QueueEvent, QueueFull, QueueManager from workflows import WorkflowBuilder @@ -218,3 +219,35 @@ async def test_prompt_enhancer_is_used(db: Database, tmp_path: Path) -> None: assert any(event.kind == "enhancing" for event in events) finally: await queue.stop() + + +async def test_blocked_prompt_is_logged(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + queue = make_queue(db, make_settings(tmp_path), events) + + with pytest.raises(PromptBlocked): + await queue.enqueue(user_id=1, mode="t2i", prompt="child nude") + + blocked = await db.list_blocked() + assert len(blocked) == 1 + assert blocked[0].user_id == 1 + assert blocked[0].reason == "csam" + assert blocked[0].prompt == "child nude" + + await queue.stop() + + +async def test_admin_prompt_is_not_filtered(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + queue = make_queue(db, make_settings(tmp_path), events) + + request, _ = await queue.enqueue( + user_id=1, + mode="t2i", + prompt="child nude", + is_admin=True, + ) + assert request.id > 0 + assert await db.list_blocked() == [] + + await queue.stop()