# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Unified Input Sanitization for Familiar

Provides consistent sanitization across all input points in the framework.

Usage:
    from familiar.core.sanitization import (
        sanitize_user_input, sanitize_for_shell, SanitizationLevel
    )

    result = sanitize_user_input(user_message, SanitizationLevel.BASIC)
    if result.was_modified:
        logger.info(f"Input sanitized: {result.warnings}")
"""

import html
import logging
import re
from dataclasses import dataclass, field
from enum import Enum
from typing import List, Set

logger = logging.getLogger(__name__)


class SanitizationLevel(str, Enum):
    """
    Levels of input sanitization strictness.

    NONE: No modification (use with caution)
    BASIC: HTML escape, null byte removal
    STRICT: Remove non-printable characters
    PARANOID: Whitelist alphanumeric + basic punctuation only
    """

    NONE = "none"
    BASIC = "basic"
    STRICT = "strict"
    PARANOID = "paranoid"


@dataclass
class SanitizationResult:
    """
    Result of a sanitization operation.

    Attributes:
        original: The original input
        sanitized: The sanitized output
        was_modified: Whether any changes were made
        warnings: List of warning messages about modifications
    """

    original: str
    sanitized: str
    was_modified: bool
    warnings: List[str] = field(default_factory=list)

    def __bool__(self) -> bool:
        """Returns True if sanitization was successful (has output)."""
        return bool(self.sanitized)


# Characters allowed in PARANOID mode
PARANOID_WHITELIST = set(
    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,!?'\"-:;@#$%&*()\n\t"
)

# Control characters to strip (except tab, newline, carriage return)
CONTROL_CHARS = set(chr(i) for i in range(32) if i not in (9, 10, 13))


def sanitize_user_input(
    text: str,
    level: SanitizationLevel = SanitizationLevel.BASIC,
    max_length: int = 10_000,
    allow_newlines: bool = True,
    allow_html: bool = False,
) -> SanitizationResult:
    """
    Sanitize user input based on security level.

    Args:
        text: Input text to sanitize
        level: Sanitization strictness level
        max_length: Maximum allowed length (truncates if exceeded)
        allow_newlines: Whether to preserve newlines in STRICT mode
        allow_html: If False, HTML-escape in BASIC+ modes

    Returns:
        SanitizationResult with sanitized text and modification info
    """
    if not isinstance(text, str):
        text = str(text)

    warnings: List[str] = []
    original = text

    # Length check (always applied)
    if len(text) > max_length:
        text = text[:max_length]
        warnings.append(f"Input truncated from {len(original)} to {max_length} characters")

    # Null byte removal (always applied - these are never valid in user input)
    if "\x00" in text:
        text = text.replace("\x00", "")
        warnings.append("Null bytes removed")

    if level == SanitizationLevel.NONE:
        pass

    elif level == SanitizationLevel.BASIC:
        # HTML escape to prevent XSS in any web rendering
        if not allow_html:
            escaped = html.escape(text)
            if escaped != text:
                text = escaped
                warnings.append("HTML entities escaped")

    elif level == SanitizationLevel.STRICT:
        # Remove control characters except newline/tab
        allowed_whitespace = {"\n", "\r", "\t"} if allow_newlines else {"\t"}
        new_text = []
        removed_count = 0

        for char in text:
            if char.isprintable() or char in allowed_whitespace:
                new_text.append(char)
            else:
                removed_count += 1

        if removed_count > 0:
            text = "".join(new_text)
            warnings.append(f"Removed {removed_count} non-printable characters")

        # Also HTML escape
        if not allow_html:
            escaped = html.escape(text)
            if escaped != text:
                text = escaped
                warnings.append("HTML entities escaped")

    elif level == SanitizationLevel.PARANOID:
        # Whitelist only: alphanumeric, spaces, basic punctuation
        new_text = []
        removed_chars: Set[str] = set()

        for char in text:
            if char in PARANOID_WHITELIST:
                new_text.append(char)
            else:
                removed_chars.add(repr(char))

        if removed_chars:
            text = "".join(new_text)
            warnings.append(
                f"Removed non-whitelisted characters: {', '.join(sorted(removed_chars)[:10])}"
                + ("..." if len(removed_chars) > 10 else "")
            )

    return SanitizationResult(
        original=original, sanitized=text, was_modified=(text != original), warnings=warnings
    )


def sanitize_for_shell(command: str) -> SanitizationResult:
    """
    Validate input intended for shell execution.

    Returns the command unchanged if safe, or marks it as modified
    with warnings if dangerous sequences are detected. Does NOT
    attempt to sanitize — shell injection cannot be fixed by replacement.
    Primary protection comes from pattern validation in pattern_safety.py.

    Args:
        command: Shell command string

    Returns:
        SanitizationResult (sanitized field = original if safe, empty if dangerous)
    """
    from .pattern_safety import check_shell_safety

    warnings: List[str] = []
    original = command

    # Remove null bytes
    if "\x00" in command:
        command = command.replace("\x00", "")
        warnings.append("Null bytes removed")

    # Validate against dangerous patterns (reject, don't replace)
    is_safe, error_msg = check_shell_safety(command)
    if not is_safe:
        warnings.append(f"Dangerous pattern detected: {error_msg}")
        return SanitizationResult(
            original=original, sanitized="", was_modified=True, warnings=warnings
        )

    return SanitizationResult(
        original=original, sanitized=command, was_modified=(command != original), warnings=warnings
    )


def sanitize_for_path(path: str) -> SanitizationResult:
    """
    Sanitize input intended for filesystem paths.

    Removes path traversal attempts and normalizes the path.

    Args:
        path: File path string

    Returns:
        SanitizationResult with sanitized path
    """

    warnings: List[str] = []
    original = path

    # Remove null bytes
    if "\x00" in path:
        path = path.replace("\x00", "")
        warnings.append("Null bytes removed")

    # Reject paths containing traversal sequences (replace is bypassable with ....///)
    traversal_patterns = ["../", "..\\", "%2e%2e/", "%2e%2e\\", ".."]
    for pattern in traversal_patterns:
        if pattern.lower() in path.lower():
            raise ValueError(f"Path traversal sequence detected: '{pattern}' in path '{original}'")

    # Remove leading slashes that would make it absolute (context-dependent)
    # Only do this if explicitly sanitizing for relative paths

    # Normalize path separators
    path = path.replace("\\", "/")

    # Remove redundant slashes
    while "//" in path:
        path = path.replace("//", "/")

    return SanitizationResult(
        original=original, sanitized=path, was_modified=(path != original), warnings=warnings
    )


def sanitize_for_logging(text: str, max_length: int = 1000) -> str:
    """
    Sanitize text for safe inclusion in logs.

    - Truncates long strings
    - Escapes control characters
    - Masks potential secrets

    Args:
        text: Text to sanitize
        max_length: Maximum length for log output

    Returns:
        Sanitized string safe for logging
    """
    if not isinstance(text, str):
        text = str(text)

    # Truncate
    if len(text) > max_length:
        text = text[:max_length] + f"... [truncated, {len(text)} total chars]"

    # Escape control characters for readability
    text = text.encode("unicode_escape").decode("ascii")

    return text


def strip_ansi_codes(text: str) -> str:
    """
    Remove ANSI escape codes from text.

    Useful for cleaning terminal output before processing.

    Args:
        text: Text potentially containing ANSI codes

    Returns:
        Clean text without ANSI codes
    """
    ansi_pattern = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]")
    return ansi_pattern.sub("", text)


# ============================================================
# HTML Allowlist Sanitization
# ============================================================

_DEFAULT_ALLOWED_TAGS: set[str] = {
    "p",
    "br",
    "h1",
    "h2",
    "h3",
    "h4",
    "h5",
    "h6",
    "ul",
    "ol",
    "li",
    "a",
    "em",
    "strong",
    "code",
    "pre",
    "blockquote",
    "table",
    "thead",
    "tbody",
    "tr",
    "td",
    "th",
    "span",
    "div",
    "b",
    "i",
    "u",
    "img",
    "figure",
    "figcaption",
    "dl",
    "dt",
    "dd",
    "hr",
    "sub",
    "sup",
    "mark",
    "abbr",
    "time",
    "details",
    "summary",
}

_DEFAULT_ALLOWED_ATTRS: dict[str, set[str]] = {
    "a": {"href"},
    "img": {"src", "alt"},
    "td": {"colspan", "rowspan"},
    "th": {"colspan", "rowspan"},
    "time": {"datetime"},
}

# Tags that are always removed (including their content), even if
# someone passes them in allowed_tags.
_DANGEROUS_TAGS: set[str] = {
    "script",
    "style",
    "iframe",
    "embed",
    "object",
    "form",
    "input",
    "textarea",
    "select",
    "button",
    "link",
    "meta",
    "base",
    "svg",
    "math",
}

_DANGEROUS_SCHEMES: set[str] = {"javascript", "data", "vbscript"}


def _is_safe_url(url: str) -> bool:
    """Check that a URL does not use a dangerous scheme."""
    normalized = url.strip().lower()
    for scheme in _DANGEROUS_SCHEMES:
        if normalized.startswith(f"{scheme}:"):
            return False
    return True


def sanitize_html(
    html_content: str,
    allowed_tags: set[str] | None = None,
    allowed_attrs: dict[str, set[str]] | None = None,
) -> str:
    """
    Allowlist-sanitize HTML using BeautifulSoup.

    Removes all tags not in *allowed_tags* (unwrapping them so text children
    are preserved).  On allowed tags, strips every attribute not listed in
    *allowed_attrs*.  ``href`` and ``src`` values are validated against
    dangerous URI schemes (``javascript:``, ``data:``, ``vbscript:``).

    Tags in ``_DANGEROUS_TAGS`` are removed entirely (including their
    content) regardless of the allowlist.

    Args:
        html_content: Raw HTML string.
        allowed_tags: Set of tag names to keep. Defaults to
            ``_DEFAULT_ALLOWED_TAGS``.
        allowed_attrs: Mapping of tag name → set of allowed attribute names.
            Defaults to ``_DEFAULT_ALLOWED_ATTRS``.

    Returns:
        Cleaned HTML string.
    """
    from bs4 import BeautifulSoup

    if allowed_tags is None:
        allowed_tags = _DEFAULT_ALLOWED_TAGS
    if allowed_attrs is None:
        allowed_attrs = _DEFAULT_ALLOWED_ATTRS

    soup = BeautifulSoup(html_content, "html.parser")

    # First pass: remove dangerous tags and all their content.
    for tag in soup.find_all(list(_DANGEROUS_TAGS)):
        tag.decompose()

    # Second pass: handle remaining tags.
    # We iterate over a copy of the list because unwrapping mutates the tree.
    for tag in list(soup.find_all(True)):
        if tag.name not in allowed_tags:
            # Not allowed — unwrap (keep text children, drop the tag itself).
            tag.unwrap()
            continue

        # Tag is allowed — strip disallowed attributes.
        tag_allowed = allowed_attrs.get(tag.name, set())
        for attr_name in list(tag.attrs):
            if attr_name not in tag_allowed:
                del tag[attr_name]

        # Validate URL-bearing attributes.
        for url_attr in ("href", "src"):
            if url_attr in tag.attrs and not _is_safe_url(tag[url_attr]):
                del tag[url_attr]

    return str(soup)


# ============================================================
# SSRF Protection
# ============================================================

import ipaddress  # noqa: E402
import socket as _socket  # noqa: E402
from urllib.parse import urlparse as _urlparse  # noqa: E402

# Private/internal IP ranges that should never be accessed via SSRF
_BLOCKED_NETWORKS = [
    ipaddress.ip_network("127.0.0.0/8"),  # Loopback
    ipaddress.ip_network("10.0.0.0/8"),  # Private Class A
    ipaddress.ip_network("172.16.0.0/12"),  # Private Class B
    ipaddress.ip_network("192.168.0.0/16"),  # Private Class C
    ipaddress.ip_network("169.254.0.0/16"),  # Link-local / cloud metadata
    ipaddress.ip_network("0.0.0.0/8"),  # "This" network
    ipaddress.ip_network("100.64.0.0/10"),  # Carrier-grade NAT
    ipaddress.ip_network("192.0.0.0/24"),  # IETF protocol assignments
    ipaddress.ip_network("198.18.0.0/15"),  # Benchmarking
    ipaddress.ip_network("::1/128"),  # IPv6 loopback
    ipaddress.ip_network("fc00::/7"),  # IPv6 unique local
    ipaddress.ip_network("fe80::/10"),  # IPv6 link-local
]

_BLOCKED_SCHEMES = {"file", "ftp", "gopher", "data", "javascript"}


def check_ssrf(url: str) -> tuple:
    """
    Validate a URL against SSRF attacks.

    Blocks requests to private/internal IP ranges, cloud metadata endpoints,
    and non-HTTP(S) schemes.

    Args:
        url: The URL to validate.

    Returns:
        (is_safe, error_message) tuple.  is_safe is True when the URL is OK.
    """
    try:
        parsed = _urlparse(url)
    except Exception:
        return False, "Invalid URL"

    # Block dangerous schemes
    scheme = (parsed.scheme or "").lower()
    if scheme in _BLOCKED_SCHEMES:
        return False, f"Blocked URL scheme: {scheme}"

    if scheme not in ("http", "https"):
        return False, f"Only http and https URLs are allowed, got: {scheme}"

    hostname = parsed.hostname
    if not hostname:
        return False, "URL has no hostname"

    # Resolve hostname to IP and check against blocked ranges
    try:
        addr_infos = _socket.getaddrinfo(hostname, None, type=_socket.SOCK_STREAM)
    except _socket.gaierror:
        return False, f"Could not resolve hostname: {hostname}"

    for family, _, _, _, sockaddr in addr_infos:
        ip_str = sockaddr[0]
        try:
            ip = ipaddress.ip_address(ip_str)
        except ValueError as e:
            logger.debug("Value conversion error suppressed: %s", e)
            continue
        for network in _BLOCKED_NETWORKS:
            if ip in network:
                return False, (
                    f"Blocked request to private/internal address: {hostname} resolves to {ip_str}"
                )

    return True, ""


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    # Enums
    "SanitizationLevel",
    # Classes
    "SanitizationResult",
    # Functions
    "sanitize_user_input",
    "sanitize_for_shell",
    "sanitize_for_path",
    "sanitize_for_logging",
    "strip_ansi_codes",
    "check_ssrf",
    "sanitize_html",
]
