gizmoduck/ingestion/guided_imports.py

"""Safe discovery of local source documents for the guided import workflow."""

import os
import stat
from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path

from django.conf import settings

from .parsers import MAX_BYTES, ParseError, parse_document


@dataclass(frozen=True)
class DiscoveredFile:
    name: str
    document_hash: str
    profile: str | None
    sections: tuple[str, ...]
    code: str | None


_PROFILES_BY_SUFFIX = {
    ".csv": ("BANK_CSV_V1",),
    ".txt": ("OUROCARD_FULL_V1", "OUROCARD_PARTIAL_V1"),
}

_SUPPORTS_SAFE_DIRECTORY_FDS = (
    os.name == "posix"
    and os.open in os.supports_dir_fd
    and os.stat in os.supports_dir_fd
    and os.listdir in os.supports_fd
    and hasattr(os, "O_DIRECTORY")
    and hasattr(os, "O_NOFOLLOW")
)


def _profile_and_sections(path: Path, data: bytes) -> tuple[str | None, tuple[str, ...]]:
    for profile in _PROFILES_BY_SUFFIX.get(path.suffix.casefold(), ()):
        try:
            result = parse_document(data, profile)
        except ParseError:
            continue
        items = tuple(item for item in result.items if not item.errors)
        if items:
            return profile, tuple(dict.fromkeys(item.section for item in items))
    return None, ()


def _unreadable(name: str) -> DiscoveredFile:
    return DiscoveredFile(name, "", None, (), "UNREADABLE_FILE")


def _close(descriptor: int | None) -> None:
    if descriptor is not None:
        try:
            os.close(descriptor)
        except OSError:
            pass


def _open_import_root() -> int | None:
    if not _SUPPORTS_SAFE_DIRECTORY_FDS:
        return None
    try:
        root = os.path.normpath(os.fspath(settings.IMPORT_ROOT))
        if not os.path.isabs(root):
            root = os.path.abspath(root)
        path = Path(root)
        descriptor = os.open(path.anchor, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
        for component in path.parts[1:]:
            next_descriptor = os.open(
                component,
                os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW,
                dir_fd=descriptor,
            )
            _close(descriptor)
            descriptor = next_descriptor
        if not stat.S_ISDIR(os.fstat(descriptor).st_mode):
            _close(descriptor)
            return None
        return descriptor
    except (OSError, RuntimeError, TypeError, ValueError):
        _close(locals().get("descriptor"))
        return None


def _read_discovered_file(name: str, root_descriptor: int, expected: os.stat_result) -> DiscoveredFile:
    descriptor = None
    try:
        flags = os.O_RDONLY | getattr(os, "O_BINARY", 0)
        if no_follow := getattr(os, "O_NOFOLLOW", 0):
            flags |= no_follow
        descriptor = os.open(name, flags, dir_fd=root_descriptor)
        opened = os.fstat(descriptor)
        if not stat.S_ISREG(opened.st_mode) or not os.path.samestat(expected, opened):
            return _unreadable(name)
        with os.fdopen(descriptor, "rb", closefd=False) as stream:
            data = stream.read(MAX_BYTES + 1)
    except OSError:
        return _unreadable(name)
    finally:
        _close(descriptor)
    if len(data) > MAX_BYTES:
        return DiscoveredFile(name, "", None, (), "INPUT_TOO_LARGE")
    document_hash = sha256(data).hexdigest()
    if Path(name).suffix.casefold() not in _PROFILES_BY_SUFFIX:
        return DiscoveredFile(name, document_hash, None, (), "UNSUPPORTED_FILE")
    profile, sections = _profile_and_sections(Path(name), data)
    if profile is None:
        return DiscoveredFile(name, document_hash, None, (), "UNRECOGNIZED_LAYOUT")
    return DiscoveredFile(name, document_hash, profile, sections, None)


def discover_imports() -> tuple[DiscoveredFile, ...]:
    """Return safe metadata for direct regular files under ``IMPORT_ROOT``."""
    root_descriptor = _open_import_root()
    if root_descriptor is None:
        return ()
    try:
        try:
            entries = os.listdir(root_descriptor)
        except OSError:
            return ()
        discovered = []
        for name in entries:
            try:
                expected = os.stat(name, dir_fd=root_descriptor, follow_symlinks=False)
            except OSError:
                discovered.append(_unreadable(name))
                continue
            if stat.S_ISREG(expected.st_mode):
                discovered.append(_read_discovered_file(name, root_descriptor, expected))
        return tuple(sorted(discovered, key=lambda row: row.name))
    finally:
        _close(root_descriptor)