omnisvg-lite/validator.py

"""Валидация SVG-ответов модели.

Делает две вещи:
1) `extract_svg()` — вытаскивает <svg>...</svg> из произвольного текста
   (модель иногда отвечает с пояснениями или code-fence).
2) `validate(svg_text, mode=...)` — строгая проверка по правилам дизайна.
3) `validate_svg(svg_text) -> (ok, reason, cleaned_svg)` — обёртка для UI,
   описанная в задаче: всегда возвращает кортеж, не бросает исключений.

Правила (соответствуют §5 design.md и §6 spec.md):
- well-formed XML (lxml);
- ровно один корневой <svg>;
- наличие `viewBox` (4 числа);
- теги — только из ALLOWED_TAGS;
- <script>, <foreignObject>, <image> — запрещены;
- http://, https:// — запрещены в href/xlink:href и в `url(...)`;
- on* event-handler атрибуты — запрещены (доп. защита по задаче);
- размер файла <= лимита mode (32/256 KB).
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass

from lxml import etree


log = logging.getLogger(__name__)


SVG_NS = "http://www.w3.org/2000/svg"

ALLOWED_TAGS: set[str] = {
    "svg",
    "g",
    "defs",
    "symbol",
    "use",
    "path",
    "rect",
    "circle",
    "ellipse",
    "line",
    "polygon",
    "polyline",
    "linearGradient",
    "radialGradient",
    "stop",
    "text",
    "tspan",
    "textPath",
    "title",
    "desc",
    "filter",
    "feGaussianBlur",
    "feOffset",
    "feMerge",
    "feMergeNode",
    "feColorMatrix",
    "feBlend",
    "feFlood",
    "feComposite",
    "clipPath",
    "mask",
    "pattern",
}

DISALLOWED_TAGS: set[str] = {"script", "foreignObject", "image", "iframe", "style"}

MAX_BYTES: dict[str, int] = {
    "icon": 32 * 1024,
    "illustration": 256 * 1024,
}


class ValidatorError(Exception):
    """Ошибка валидации с машино-читаемым кодом и человеческим пояснением."""

    def __init__(self, code: str, detail: str = ""):
        super().__init__(f"{code}: {detail}" if detail else code)
        self.code = code
        self.detail = detail


_SVG_RE = re.compile(
    r"<svg\b[^>]*>.*?</svg\s*>",
    re.IGNORECASE | re.DOTALL,
)


def extract_svg(text: str) -> str:
    """Достаёт первое вхождение <svg>...</svg> из произвольного ответа модели."""
    if not text or "<svg" not in text.lower():
        raise ValidatorError("not_svg", "в ответе нет <svg>...</svg>")
    match = _SVG_RE.search(text)
    if not match:
        raise ValidatorError("not_svg", "не удалось выделить <svg>...</svg> regex'ом")
    return match.group(0).strip()


def _local(tag: object) -> str:
    """Снимает namespace: {ns}local -> local. Терпимо к non-str (cyfunction)."""
    s = str(tag)
    if "}" in s:
        return s.split("}", 1)[1]
    return s


def _iter_real_elements(root: etree._Element):
    """Итератор только по Element-узлам (пропускает комменты и PI)."""
    for el in root.iter():
        # lxml._Element и etree._Element имеют tag как str; cyfunction — это
        # callable у комментариев/PI. Отфильтруем через isinstance.
        if isinstance(getattr(el, "tag", None), str):
            yield el


def _check_viewbox(root: etree._Element) -> None:
    vb = root.get("viewBox")
    if not vb:
        raise ValidatorError("missing_viewbox", "атрибут viewBox отсутствует")
    try:
        parts = [float(x) for x in vb.split()]
    except ValueError as exc:
        raise ValidatorError("bad_viewbox", f"viewBox='{vb}'") from exc
    if len(parts) != 4:
        raise ValidatorError("bad_viewbox", f"viewBox='{vb}' (нужно 4 числа)")


def _check_tags_and_attrs(root: etree._Element) -> None:
    for el in _iter_real_elements(root):
        local = _local(el.tag)
        if local in DISALLOWED_TAGS:
            raise ValidatorError("disallowed_tag", local)
        if local not in ALLOWED_TAGS:
            raise ValidatorError("unknown_tag", local)
        for attr, val in el.attrib.items():
            if not val:
                continue
            attr_local = _local(attr).lower()
            if attr_local.startswith("on"):
                raise ValidatorError("event_handler", f"{attr_local}={val[:80]}")
            if "http://" in val or "https://" in val:
                if attr_local.endswith("href") or "url(" in val.lower():
                    raise ValidatorError("external_ref", f"{attr_local}={val[:80]}")


def _ensure_size(svg_text: str, mode: str) -> None:
    limit = MAX_BYTES.get(mode)
    if limit is None:
        return
    if len(svg_text.encode("utf-8")) > limit:
        raise ValidatorError(
            "too_large",
            f"{len(svg_text.encode('utf-8'))} байт > {limit} байт (mode={mode})",
        )


def _serialize_clean(root: etree._Element) -> str:
    if not root.get("xmlns"):
        new_root = etree.Element(_local(root.tag), nsmap={None: SVG_NS})
        for k, v in root.attrib.items():
            new_root.set(_local(k), v)
        for child in root:
            new_root.append(child)
        root = new_root
    xml_bytes = etree.tostring(
        root, pretty_print=False, xml_declaration=False, encoding="utf-8"
    )
    return xml_bytes.decode("utf-8")


def validate(svg_text: str, *, mode: str = "icon") -> etree._Element:
    """Строгая валидация по правилам дизайна. Бросает ValidatorError."""
    svg_text = extract_svg(svg_text)
    try:
        root = etree.fromstring(svg_text.encode("utf-8"))
    except etree.XMLSyntaxError as exc:
        raise ValidatorError("malformed_xml", str(exc)) from exc

    if _local(root.tag) != "svg":
        raise ValidatorError("not_svg", root.tag)

    _check_viewbox(root)
    _check_tags_and_attrs(root)
    _ensure_size(svg_text, mode)
    return root


@dataclass(frozen=True)
class ValidationResult:
    ok: bool
    reason: str
    cleaned_svg: str


def validate_svg(svg_text: str, *, mode: str = "icon") -> tuple[bool, str, str]:
    """Удобная обёртка над `validate()`: всегда возвращает кортеж.

    Returns:
        Кортеж `(ok, reason, cleaned_svg)`.
    """
    if not svg_text or not svg_text.strip():
        return False, "empty_input", ""
    try:
        root = validate(svg_text, mode=mode)
    except ValidatorError as exc:
        log.info("validate_svg: reject (%s)", exc.code)
        return False, exc.code, ""
    except Exception as exc:  # noqa: BLE001
        log.exception("validate_svg: неожиданная ошибка")
        return False, f"internal:{type(exc).__name__}", ""
    return True, "", _serialize_clean(root)


__all__ = [
    "ALLOWED_TAGS",
    "DISALLOWED_TAGS",
    "MAX_BYTES",
    "SVG_NS",
    "ValidationResult",
    "ValidatorError",
    "extract_svg",
    "validate",
    "validate_svg",
]