n0mad1k
142 lines
4.3 KiB
Python
142 lines
4.3 KiB
Python
"""Input sanitization for Mosaic — filenames, paths, content types."""
|
|
import os
|
|
import re
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Allowed characters in filenames
|
|
FILENAME_SAFE = re.compile(r'[^a-zA-Z0-9._-]')
|
|
# Maximum filename length
|
|
MAX_FILENAME_LENGTH = 200
|
|
# ANSI escape sequence pattern
|
|
ANSI_ESCAPE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07')
|
|
|
|
# Magic bytes → MIME type mapping
|
|
MAGIC_BYTES = {
|
|
b'%PDF': 'application/pdf',
|
|
b'\x89PNG': 'image/png',
|
|
b'\xff\xd8\xff': 'image/jpeg',
|
|
b'GIF87a': 'image/gif',
|
|
b'GIF89a': 'image/gif',
|
|
b'PK\x03\x04': 'application/zip',
|
|
b'<!DOCTYPE': 'text/html',
|
|
b'<html': 'text/html',
|
|
b'<HTML': 'text/html',
|
|
}
|
|
|
|
|
|
def sanitize_filename(filename: str) -> str:
|
|
"""Sanitize a filename to safe characters only.
|
|
|
|
- Strip to basename (no directory traversal)
|
|
- Replace unsafe characters with underscore
|
|
- Limit length
|
|
"""
|
|
# Strip to basename
|
|
filename = os.path.basename(filename)
|
|
# Replace unsafe characters
|
|
sanitized = FILENAME_SAFE.sub('_', filename)
|
|
# Collapse multiple underscores
|
|
sanitized = re.sub(r'_+', '_', sanitized)
|
|
# Limit length
|
|
if len(sanitized) > MAX_FILENAME_LENGTH:
|
|
name, ext = os.path.splitext(sanitized)
|
|
sanitized = name[:MAX_FILENAME_LENGTH - len(ext)] + ext
|
|
# Ensure non-empty
|
|
if not sanitized or sanitized == '_':
|
|
sanitized = 'unnamed'
|
|
return sanitized
|
|
|
|
|
|
def sanitize_path(path: str, base_dir: str) -> Optional[str]:
|
|
"""Validate a path stays within base_dir (path traversal prevention).
|
|
|
|
Returns the resolved path if safe, None if path escapes base_dir.
|
|
"""
|
|
base = os.path.realpath(base_dir)
|
|
resolved = os.path.realpath(os.path.join(base, path))
|
|
if resolved.startswith(base + os.sep) or resolved == base:
|
|
return resolved
|
|
logger.warning("Path traversal blocked: %s → %s (base: %s)", path, resolved, base)
|
|
return None
|
|
|
|
|
|
def validate_content_type(file_path: str, expected_type: Optional[str] = None) -> str:
|
|
"""Determine content type by magic bytes, not file extension.
|
|
|
|
Args:
|
|
file_path: Path to the file
|
|
expected_type: Expected MIME type to validate against
|
|
|
|
Returns:
|
|
Detected MIME type string
|
|
|
|
Raises:
|
|
ValueError: If expected_type is provided and doesn't match
|
|
"""
|
|
try:
|
|
import magic
|
|
detected = magic.from_file(file_path, mime=True)
|
|
except ImportError:
|
|
# Fallback to magic bytes check
|
|
detected = _detect_by_magic_bytes(file_path)
|
|
|
|
if expected_type and detected != expected_type:
|
|
# Allow some flexibility (e.g., text/plain for text/html)
|
|
if not _types_compatible(detected, expected_type):
|
|
raise ValueError(
|
|
f"Content type mismatch: expected {expected_type}, detected {detected}"
|
|
)
|
|
return detected
|
|
|
|
|
|
def _detect_by_magic_bytes(file_path: str) -> str:
|
|
"""Fallback detection using magic bytes."""
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
header = f.read(16)
|
|
except (IOError, OSError):
|
|
return 'application/octet-stream'
|
|
|
|
for magic, mime in MAGIC_BYTES.items():
|
|
if header.startswith(magic):
|
|
return mime
|
|
|
|
# Check if it looks like text
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
f.read(1024)
|
|
return 'text/plain'
|
|
except (UnicodeDecodeError, IOError):
|
|
return 'application/octet-stream'
|
|
|
|
|
|
def _types_compatible(detected: str, expected: str) -> bool:
|
|
"""Check if two MIME types are compatible."""
|
|
# Same type
|
|
if detected == expected:
|
|
return True
|
|
# text/plain is compatible with text/*
|
|
if detected == 'text/plain' and expected.startswith('text/'):
|
|
return True
|
|
# text/html variants
|
|
if detected.startswith('text/html') and expected.startswith('text/html'):
|
|
return True
|
|
return False
|
|
|
|
|
|
def strip_ansi(text: str) -> str:
|
|
"""Strip ANSI escape sequences from text."""
|
|
return ANSI_ESCAPE.sub('', text)
|
|
|
|
|
|
def sanitize_for_display(text: str, max_length: int = 0) -> str:
|
|
"""Sanitize text for safe Rich display — strip ANSI, optionally truncate."""
|
|
clean = strip_ansi(text)
|
|
if max_length > 0 and len(clean) > max_length:
|
|
clean = clean[:max_length] + "..."
|
|
return clean
|