dlp-io/dlp_io.py

"""io-compatible file reads through an approved Python helper on Windows DLP hosts."""

from __future__ import annotations

import atexit
import builtins
import io
import json
import operator
import os
import queue
import shutil
import subprocess
import threading
from dataclasses import dataclass
from enum import Enum
from multiprocessing.context import AuthenticationError
from multiprocessing.connection import Client

__version__ = "0.1.1"

__all__ = [
    "DlpConfig",
    "DlpConfigurationError",
    "DlpHelperStartError",
    "DlpIoError",
    "DlpProtocolError",
    "DlpSession",
    "DlpSessionBusyError",
    "DlpTransportError",
    "configure",
    "install_open_patch",
    "install_reader_patch",
    "open",
    "shutdown",
    "uninstall_open_patch",
]


# ---------------------------------------------------------------------------
# Errors


class DlpIoError(Exception):
    """Base exception for dlp-io failures."""


class DlpConfigurationError(DlpIoError):
    """The bridge configuration is invalid for the current lifecycle state."""


class DlpHelperStartError(DlpIoError):
    """The approved Python helper could not be started."""


class DlpProtocolError(DlpIoError):
    """The helper and client exchanged an invalid protocol message."""


class DlpTransportError(DlpIoError):
    """The authenticated pipe stream failed before a valid EOF."""


class DlpSessionBusyError(DlpIoError):
    """A session already owns an active file stream."""


# ---------------------------------------------------------------------------
# Configuration


@dataclass(frozen=True)
class DlpConfig:
    python_executable: str | None = None
    startup_timeout: float = 60

    def python_command(self) -> list[str] | None:
        if self.python_executable is None:
            return None
        return [self.python_executable]


# ---------------------------------------------------------------------------
# Helper protocol

PROTOCOL_VERSION = 1
CONTROL_LIMIT = 64 * 1024
CHUNK_SIZE = 1024 * 1024
HANDSHAKE_PREFIX = "DLP_IO 1 "


class HelperErrorCode(str, Enum):
    INVALID_PATH = "invalid_path"
    INVALID_OFFSET = "invalid_offset"
    FILE_ACCESS = "file_access"
    INVALID_REQUEST = "invalid_request"
    INVALID_REQUEST_TYPE = "invalid_request_type"
    UNSUPPORTED_VERSION = "unsupported_version"
    UNKNOWN_OPERATION = "unknown_operation"
    INTERNAL = "internal"


@dataclass(frozen=True)
class HelperSourceConfig:
    authkey: bytes
    protocol_version: int = PROTOCOL_VERSION
    control_limit: int = CONTROL_LIMIT
    chunk_size: int = CHUNK_SIZE

    @classmethod
    def from_authkey(cls, authkey: bytes) -> HelperSourceConfig:
        if not isinstance(authkey, bytes) or len(authkey) != 32:
            raise DlpConfigurationError("authkey must contain exactly 32 bytes")
        return cls(authkey=authkey)


_HELPER_TEMPLATE = r'''
import json
import os
import sys
from multiprocessing.connection import AuthenticationError, Listener

VERSION = __PROTOCOL_VERSION__
CONTROL_LIMIT = __CONTROL_LIMIT__
CHUNK_SIZE = __CHUNK_SIZE__
PIPE_PREFIX = r"\\.\pipe\DlpIo_"
AUTHKEY = bytes.fromhex("__AUTHKEY_HEX__")


def send_json(conn, value):
    conn.send_bytes(json.dumps(value).encode("utf-8"))


def error(conn, code, message):
    send_json(conn, {
        "version": VERSION,
        "ok": False,
        "error_code": code,
        "error": message,
    })


def handle_read(conn, request):
    path = request.get("path")
    offset = request.get("offset", 0)
    if not isinstance(path, str) or not path:
        error(conn, "invalid_path", "path must be a non-empty string")
        return
    if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0:
        error(conn, "invalid_offset", "offset must be a non-negative integer")
        return
    try:
        size = os.path.getsize(path)
        stream = open(path, "rb")
    except Exception as exc:
        error(conn, "file_access", "%s: %s" % (type(exc).__name__, exc))
        return
    try:
        stream.seek(offset)
        send_json(conn, {"version": VERSION, "ok": True, "size": size})
        while True:
            chunk = stream.read(CHUNK_SIZE)
            if not chunk:
                conn.send_bytes(b"")
                return
            conn.send_bytes(chunk)
    except (BrokenPipeError, EOFError, OSError):
        return
    finally:
        stream.close()


def handle(conn):
    try:
        request = json.loads(conn.recv_bytes(maxlength=CONTROL_LIMIT).decode("utf-8"))
    except Exception as exc:
        error(conn, "invalid_request", "invalid request: %s" % exc)
        return True
    if not isinstance(request, dict):
        error(conn, "invalid_request_type", "request must be a JSON object")
        return True
    if request.get("version") != VERSION:
        error(conn, "unsupported_version", "unsupported protocol version")
        return True
    operation = request.get("op")
    if operation == "ping":
        send_json(conn, {
            "version": VERSION,
            "ok": True,
            "pid": os.getpid(),
            "python": sys.version,
        })
        return True
    if operation == "read":
        handle_read(conn, request)
        return True
    if operation == "quit":
        send_json(conn, {"version": VERSION, "ok": True})
        return False
    error(conn, "unknown_operation", "unknown operation: %r" % operation)
    return True


def main():
    address = PIPE_PREFIX + os.urandom(16).hex()
    listener = Listener(address=address, family="AF_PIPE", authkey=AUTHKEY)
    sys.stdout.write("DLP_IO " + str(VERSION) + " " + address + "\n")
    sys.stdout.flush()
    running = True
    while running:
        try:
            conn = listener.accept()
        except AuthenticationError:
            continue
        try:
            running = handle(conn)
        except Exception as exc:
            try:
                error(conn, "internal", "%s: %s" % (type(exc).__name__, exc))
            except Exception:
                pass
        finally:
            try:
                conn.close()
            except Exception:
                pass
    listener.close()


try:
    main()
except Exception as exc:
    sys.stderr.write("dlp-io helper fatal: %s: %s\n" % (type(exc).__name__, exc))
    sys.stderr.flush()
    sys.exit(1)
'''


def render_helper_source(authkey: bytes) -> str:
    config = HelperSourceConfig.from_authkey(authkey)
    replacements = {
        "__AUTHKEY_HEX__": config.authkey.hex(),
        "__PROTOCOL_VERSION__": str(config.protocol_version),
        "__CONTROL_LIMIT__": str(config.control_limit),
        "__CHUNK_SIZE__": str(config.chunk_size),
    }
    source = _HELPER_TEMPLATE
    for marker, value in replacements.items():
        source = source.replace(marker, value)
    return source


# ---------------------------------------------------------------------------
# Pipe-backed raw stream


class _PipeRawIO(io.RawIOBase):
    _DISCARD_LIMIT = 64 * 1024 * 1024

    def __init__(self, *, conn, size, path, offset, reopen, release) -> None:
        super().__init__()
        self._conn = conn
        self._size = int(size)
        self._path = path
        self._pos = int(offset)
        self._expected_end = max(self._size, self._pos)
        self._reopen = reopen
        self._release = release
        self._buffer = bytearray()
        self._eof = False
        self._released = False

    def readable(self) -> bool:
        return True

    def seekable(self) -> bool:
        return True

    def tell(self) -> int:
        if self.closed:
            raise ValueError("I/O operation on closed file")
        return self._pos

    def _receive(self) -> None:
        try:
            chunk = self._conn.recv_bytes(maxlength=CHUNK_SIZE)
        except (EOFError, OSError) as exc:
            raise DlpTransportError("helper disconnected before EOF") from exc
        if chunk:
            received_end = self._pos + len(self._buffer) + len(chunk)
            if received_end > self._expected_end:
                raise DlpTransportError(
                    "file size changed during transfer: expected end %d, received at least %d"
                    % (self._expected_end, received_end)
                )
            self._buffer.extend(chunk)
            return
        if self._pos != self._expected_end:
            raise DlpTransportError(
                "file size changed during transfer: expected end %d, received %d"
                % (self._expected_end, self._pos)
            )
        self._eof = True

    def readinto(self, target) -> int:
        if self.closed:
            raise ValueError("I/O operation on closed file")
        if len(target) == 0:
            return 0
        if not self._buffer:
            if self._eof:
                return 0
            self._receive()
            if self._eof:
                return 0
        count = min(len(target), len(self._buffer))
        target[:count] = self._buffer[:count]
        del self._buffer[:count]
        self._pos += count
        return count

    def seek(self, offset, whence=io.SEEK_SET) -> int:
        if self.closed:
            raise ValueError("I/O operation on closed file")
        if whence == io.SEEK_SET:
            target = offset
        elif whence == io.SEEK_CUR:
            target = self._pos + offset
        elif whence == io.SEEK_END:
            target = self._size + offset
        else:
            raise ValueError("invalid whence: %r" % (whence,))
        if target < 0:
            raise ValueError("negative seek position: %d" % target)
        delta = target - self._pos
        if delta == 0:
            return self._pos
        if 0 < delta <= self._DISCARD_LIMIT:
            remaining = delta
            while remaining:
                scratch = bytearray(min(remaining, CHUNK_SIZE))
                count = self.readinto(scratch)
                if count == 0:
                    self._pos = target
                    self._expected_end = max(self._size, target)
                    break
                remaining -= count
            return self._pos

        self._conn.close()
        try:
            conn, size = self._reopen(self._path, target)
        except Exception:
            self.close()
            raise
        if int(size) != self._size:
            conn.close()
            self.close()
            raise DlpTransportError(
                "file size changed during seek: initial %d, reopened %d"
                % (self._size, size)
            )
        self._conn = conn
        self._pos = target
        self._expected_end = max(self._size, target)
        self._buffer.clear()
        self._eof = False
        return self._pos

    def _release_once(self) -> None:
        if self._released:
            return
        self._released = True
        self._release(self)

    def close(self) -> None:
        if not self.closed:
            try:
                self._conn.close()
            finally:
                self._release_once()
        super().close()


# ---------------------------------------------------------------------------
# Helper session

_SHUTDOWN_TIMEOUT = 5


def _locate_python() -> list[str] | None:
    for variable in ("DLP_IO_PYTHON", "DATAPACKER_DLP_PYTHON"):
        value = os.environ.get(variable)
        if value:
            return [value]
    executable = shutil.which("python")
    if executable:
        return [executable]
    if shutil.which("py"):
        return ["py", "-3"]
    return None


def _helper_process_options(os_name=None) -> dict:
    if (os_name or os.name) == "nt":
        return {"creationflags": subprocess.CREATE_NO_WINDOW}
    return {}


def _stop_process(process) -> None:
    if process.poll() is not None:
        return
    process.terminate()
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        process.kill()
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            pass


def _close_process_pipes(process) -> None:
    for name in ("stdin", "stdout", "stderr"):
        stream = getattr(process, name, None)
        if stream is None:
            continue
        try:
            stream.close()
        except OSError:
            pass


class DlpSession:
    """One authenticated, short-lived approved-Python helper session."""

    def __init__(self, process, pipe_name: str, authkey: bytes) -> None:
        self._proc = process
        self.pipe_name = pipe_name
        self._authkey = authkey
        self._request_lock = threading.Lock()
        self._stream_lock = threading.Lock()
        self._state_lock = threading.RLock()
        self._active_raw = None
        self._closed = False

    @classmethod
    def start(cls, *, python_cmd=None, timeout=60):
        if os.name != "nt":
            raise DlpConfigurationError("dlp-io bridge reads require Windows AF_PIPE")
        if timeout <= 0:
            raise DlpConfigurationError("startup timeout must be positive")
        command = list(python_cmd) if python_cmd else _locate_python()
        if not command:
            raise DlpHelperStartError(
                "no approved Python found; configure DLP_IO_PYTHON"
            )
        authkey = os.urandom(32)
        source = render_helper_source(authkey)
        try:
            process = subprocess.Popen(
                command + ["-u", "-"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                **_helper_process_options(),
            )
        except OSError as exc:
            raise DlpHelperStartError(
                "failed to start approved Python %r: %s" % (command, exc)
            ) from exc
        try:
            process.stdin.write(source.encode("utf-8"))
            process.stdin.close()
        except (BrokenPipeError, OSError) as exc:
            _stop_process(process)
            raise DlpHelperStartError("failed to inject helper source: %s" % exc) from exc

        result = queue.Queue()

        def read_handshake() -> None:
            try:
                result.put(process.stdout.readline())
            except Exception as exc:
                result.put(exc)

        thread = threading.Thread(target=read_handshake, daemon=True)
        thread.start()
        try:
            line = result.get(timeout=timeout)
        except queue.Empty as exc:
            _stop_process(process)
            raise DlpHelperStartError(
                "timed out waiting for helper handshake after %s seconds" % timeout
            ) from exc
        if isinstance(line, Exception):
            _stop_process(process)
            raise DlpHelperStartError("failed to read helper handshake: %s" % line)
        decoded = line.decode("utf-8", "replace").strip()
        if not decoded.startswith(HANDSHAKE_PREFIX):
            _stop_process(process)
            stderr = process.stderr.read().decode("utf-8", "replace").strip()
            raise DlpHelperStartError(
                "invalid helper handshake %r; stderr: %s" % (decoded, stderr)
            )
        return cls(process, decoded[len(HANDSHAKE_PREFIX):], authkey)

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("dlp-io session is closed")

    def _validate_header(self, header) -> dict:
        if not isinstance(header, dict):
            raise DlpProtocolError("helper response must be a JSON object")
        if header.get("version") != PROTOCOL_VERSION:
            raise DlpProtocolError("helper response has an unsupported version")
        if not isinstance(header.get("ok"), bool):
            raise DlpProtocolError("helper response is missing boolean ok")
        return header

    def _request(self, request):
        with self._request_lock:
            try:
                encoded = json.dumps(request).encode("utf-8")
            except (TypeError, ValueError) as exc:
                raise DlpProtocolError("helper request is not JSON serializable") from exc
            try:
                conn = Client(self.pipe_name, family="AF_PIPE", authkey=self._authkey)
            except AuthenticationError as exc:
                raise DlpTransportError("failed to authenticate helper pipe") from exc
            except (EOFError, OSError) as exc:
                raise DlpTransportError("failed to connect to helper pipe") from exc
            try:
                try:
                    conn.send_bytes(encoded)
                except (EOFError, OSError) as exc:
                    raise DlpTransportError("failed to send helper request") from exc
                try:
                    payload = conn.recv_bytes(maxlength=CONTROL_LIMIT)
                except (EOFError, OSError) as exc:
                    raise DlpTransportError("failed to receive helper response") from exc
                try:
                    header = json.loads(payload.decode("utf-8"))
                except (UnicodeDecodeError, json.JSONDecodeError) as exc:
                    raise DlpProtocolError("helper response is not valid UTF-8 JSON") from exc
                return conn, self._validate_header(header)
            except Exception:
                conn.close()
                raise

    def ping(self) -> dict:
        with self._state_lock:
            self._ensure_open()
        conn, header = self._request({"version": PROTOCOL_VERSION, "op": "ping"})
        conn.close()
        if not header["ok"]:
            raise DlpProtocolError(header.get("error", "ping failed"))
        return header

    def _open_conn(self, path: str, offset: int = 0):
        conn, header = self._request(
            {
                "version": PROTOCOL_VERSION,
                "op": "read",
                "path": path,
                "offset": offset,
            }
        )
        if not header["ok"]:
            conn.close()
            raise OSError(header.get("error", "helper read failed"))
        size = header.get("size")
        if isinstance(size, bool) or not isinstance(size, int) or size < 0:
            conn.close()
            raise DlpProtocolError("helper returned an invalid file size")
        return conn, size

    def _release_stream(self, raw) -> None:
        with self._state_lock:
            if self._active_raw is raw:
                self._active_raw = None
                self._stream_lock.release()

    def _reopen_stream(self, path: str, offset: int):
        with self._state_lock:
            self._ensure_open()
        return self._open_conn(path, offset)

    def open_raw(self, path):
        with self._state_lock:
            self._ensure_open()
        if not self._stream_lock.acquire(blocking=False):
            raise DlpSessionBusyError("dlp-io session already has an active stream")
        normalized = os.fsdecode(os.fspath(path))
        try:
            conn, size = self._open_conn(normalized)
            raw = _PipeRawIO(
                conn=conn,
                size=size,
                path=normalized,
                offset=0,
                reopen=self._reopen_stream,
                release=self._release_stream,
            )
            with self._state_lock:
                if self._closed:
                    raw.close()
                    raise RuntimeError("dlp-io session is closed")
                self._active_raw = raw
            return raw
        except Exception:
            self._stream_lock.release()
            raise

    def open(
        self,
        file,
        mode="r",
        buffering=-1,
        encoding=None,
        errors=None,
        newline=None,
    ):
        buffering = _validate_path_mode(mode, buffering, encoding, errors, newline)
        if "r" not in mode or "+" in mode:
            raise io.UnsupportedOperation("DlpSession.open supports read-only modes")

        return _wrap_reader(
            self.open_raw(file), mode, buffering, encoding, errors, newline
        )

    def open_read(self, path):
        """Compatibility alias returning a buffered binary reader."""
        return self.open(path, "rb")

    def close(self) -> None:
        with self._state_lock:
            if self._closed:
                return
            self._closed = True
            active = self._active_raw
        if active is not None:
            active.close()

        def request_quit() -> None:
            try:
                conn, header = self._request(
                    {"version": PROTOCOL_VERSION, "op": "quit"}
                )
                conn.close()
                if not header["ok"]:
                    raise DlpProtocolError(header.get("error", "helper quit failed"))
            except Exception:
                pass

        quit_thread = threading.Thread(target=request_quit, daemon=True)
        quit_thread.start()
        quit_thread.join(timeout=_SHUTDOWN_TIMEOUT)
        try:
            if quit_thread.is_alive():
                _stop_process(self._proc)
                quit_thread.join(timeout=_SHUTDOWN_TIMEOUT)
            else:
                try:
                    self._proc.wait(timeout=_SHUTDOWN_TIMEOUT)
                except subprocess.TimeoutExpired:
                    _stop_process(self._proc)
        finally:
            _close_process_pipes(self._proc)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()
        return False


# ---------------------------------------------------------------------------
# Public io-compatible API

_ORIGINAL_OPEN = builtins.open
_DEFAULT_BUFFER_SIZE = 1024 * 1024
_state_lock = threading.RLock()
_config = DlpConfig()
_default_session = None


def _validate_path_mode(mode, buffering, encoding, errors, newline) -> int:
    buffering = operator.index(buffering)
    if buffering < -1:
        raise ValueError("invalid buffering size")
    if not isinstance(mode, str):
        raise TypeError("open() argument 'mode' must be str, not %s" % type(mode).__name__)
    allowed = set("axrwb+t")
    invalid = set(mode) - allowed
    duplicate = any(mode.count(character) > 1 for character in allowed)
    if not mode or invalid or duplicate or sum(c in mode for c in "axrw") != 1:
        raise ValueError("invalid mode: %r" % mode)
    if "b" in mode and "t" in mode:
        raise ValueError("can't have text and binary mode at once")
    if "b" in mode:
        if encoding is not None:
            raise ValueError("binary mode doesn't take an encoding argument")
        if errors is not None:
            raise ValueError("binary mode doesn't take an errors argument")
        if newline is not None:
            raise ValueError("binary mode doesn't take a newline argument")
    elif buffering == 0:
        raise ValueError("can't have unbuffered text I/O")
    return buffering


def configure(*, python_executable=None, startup_timeout=60) -> None:
    """Configure the lazy default session before its first bridged read."""
    global _config
    with _state_lock:
        if _default_session is not None:
            raise DlpConfigurationError(
                "cannot configure dlp_io while the default session is active"
            )
        if startup_timeout <= 0:
            raise DlpConfigurationError("startup_timeout must be positive")
        _config = DlpConfig(
            python_executable=python_executable,
            startup_timeout=startup_timeout,
        )


def _get_default_session():
    global _default_session
    with _state_lock:
        if _default_session is None:
            _default_session = DlpSession.start(
                python_cmd=_config.python_command(),
                timeout=_config.startup_timeout,
            )
        return _default_session


def _wrap_reader(raw, mode, buffering, encoding, errors, newline):
    binary = "b" in mode
    if binary and buffering == 0:
        return raw
    if not binary and buffering == 0:
        raw.close()
        raise ValueError("can't have unbuffered text I/O")
    buffer_size = _DEFAULT_BUFFER_SIZE if buffering in (-1, 1) else buffering
    if buffer_size <= 0:
        raw.close()
        raise ValueError("invalid buffering size")
    buffered = io.BufferedReader(raw, buffer_size=buffer_size)
    if binary:
        return buffered
    return io.TextIOWrapper(
        buffered,
        encoding=encoding,
        errors=errors,
        newline=newline,
        line_buffering=buffering == 1,
    )


def open(
    file,
    mode="r",
    buffering=-1,
    encoding=None,
    errors=None,
    newline=None,
    closefd=True,
    opener=None,
):
    """Open a file with io.open-compatible calling conventions."""
    is_path = isinstance(file, (str, bytes, os.PathLike))
    if not is_path or opener is not None:
        return _ORIGINAL_OPEN(
            file,
            mode,
            buffering,
            encoding,
            errors,
            newline,
            closefd,
            opener,
        )
    buffering = _validate_path_mode(mode, buffering, encoding, errors, newline)
    if "+" in mode:
        raise io.UnsupportedOperation(
            "dlp_io does not support update mode for path-based files"
        )
    pure_read = "r" in mode
    if pure_read:
        if not closefd:
            raise ValueError("Cannot use closefd=False with file name")
        path = os.fsdecode(os.fspath(file))
        raw = _get_default_session().open_raw(path)
        return _wrap_reader(raw, mode, buffering, encoding, errors, newline)
    return _ORIGINAL_OPEN(file, mode, buffering, encoding, errors, newline, closefd)


def shutdown() -> None:
    """Close and forget the lazy default session. This operation is idempotent."""
    global _default_session
    with _state_lock:
        session = _default_session
        _default_session = None
    if session is not None:
        session.close()


atexit.register(shutdown)


# ---------------------------------------------------------------------------
# Global open() patch

_ORIGINAL_IO_OPEN = io.open
_patch_lock = threading.RLock()
_installed = False


def _reader_wrapper(open_read_func):
    def patched_open(
        file,
        mode="r",
        buffering=-1,
        encoding=None,
        errors=None,
        newline=None,
        closefd=True,
        opener=None,
    ):
        is_path = isinstance(file, (str, bytes, os.PathLike))
        if not is_path or opener is not None:
            return _ORIGINAL_OPEN(
                file,
                mode,
                buffering,
                encoding,
                errors,
                newline,
                closefd,
                opener,
            )
        buffering = _validate_path_mode(mode, buffering, encoding, errors, newline)
        if "+" in mode:
            raise io.UnsupportedOperation(
                "dlp_io does not support update mode for path-based files"
            )
        pure_read = "r" in mode
        if pure_read:
            if not closefd:
                raise ValueError("Cannot use closefd=False with file name")
            reader = open_read_func(os.fsdecode(os.fspath(file)))
            if isinstance(reader, io.RawIOBase):
                return _wrap_reader(reader, mode, buffering, encoding, errors, newline)
            if "b" in mode:
                return reader
            return io.TextIOWrapper(
                reader,
                encoding=encoding,
                errors=errors,
                newline=newline,
                line_buffering=buffering == 1,
            )
        return _ORIGINAL_OPEN(file, mode, buffering, encoding, errors, newline, closefd)

    return patched_open


def install_open_patch() -> None:
    """Route global path-based reads through the lazy default dlp-io session."""
    global _installed
    with _patch_lock:
        if _installed:
            return
        builtins.open = open
        io.open = open
        _installed = True


def install_reader_patch(open_read_func) -> None:
    """Compatibility hook for an explicitly managed session.open_raw callback."""
    global _installed
    with _patch_lock:
        if _installed:
            return
        wrapper = _reader_wrapper(open_read_func)
        builtins.open = wrapper
        io.open = wrapper
        _installed = True


def uninstall_open_patch() -> None:
    global _installed
    with _patch_lock:
        if not _installed:
            return
        builtins.open = _ORIGINAL_OPEN
        io.open = _ORIGINAL_IO_OPEN
        _installed = False