MpakoImageAIbot/moderation.py

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