pipeline/validation.py

"""
Input validation — rejects malformed requests before any heavy work.

Uses cores.vision.sniff_format for magic-byte validation — no duplicated
image-signature table.
"""

from __future__ import annotations

import base64
from dataclasses import dataclass
from typing import Optional
from urllib.parse import urlparse

from cores.vision import sniff_format


@dataclass
class ValidationResult:
    valid: bool
    error: Optional[str] = None
    image_bytes: Optional[bytes] = None
    source: str = ""
    format: Optional[str] = None


class InputValidator:
    """Validates inbound image input."""

    def __init__(self, max_bytes: int = 20 * 1024 * 1024) -> None:
        self._max_bytes = max_bytes

    def validate(
        self,
        image_url: Optional[str] = None,
        image_base64: Optional[str] = None,
        image_bytes: Optional[bytes] = None,
    ) -> ValidationResult:
        if not any([image_url, image_base64, image_bytes]):
            return ValidationResult(False, error="No image input provided.")

        # URL
        if image_url:
            parsed = urlparse(image_url)
            if parsed.scheme not in ("http", "https"):
                return ValidationResult(False, error=f"Unsupported URL scheme: {parsed.scheme}")
            if not parsed.netloc:
                return ValidationResult(False, error="URL missing host.")
            host = parsed.netloc.split(":")[0].lower()
            if host in ("localhost", "127.0.0.1", "0.0.0.0", "::1"):
                return ValidationResult(False, error="Localhost URLs not permitted.")
            return ValidationResult(True, source="url")

        # Base64
        if image_base64:
            try:
                raw = image_base64.split(",", 1)[-1]
                decoded = base64.b64decode(raw, validate=True)
            except Exception as e:
                return ValidationResult(False, error=f"Invalid base64: {e}")
            if len(decoded) > self._max_bytes:
                return ValidationResult(False, error=f"Decoded image exceeds {self._max_bytes} bytes")
            fmt = sniff_format(decoded)
            if fmt is None:
                return ValidationResult(False, error="Unrecognized image format (magic bytes mismatch).")
            return ValidationResult(True, image_bytes=decoded, source="base64", format=fmt)

        # Raw bytes
        if image_bytes:
            if len(image_bytes) > self._max_bytes:
                return ValidationResult(False, error=f"Image exceeds {self._max_bytes} bytes")
            fmt = sniff_format(image_bytes)
            if fmt is None:
                return ValidationResult(False, error="Unrecognized image format (magic bytes mismatch).")
            return ValidationResult(True, image_bytes=image_bytes, source="bytes", format=fmt)

        return ValidationResult(False, error="Unreachable.")