from __future__ import annotations from dataclasses import dataclass from enum import Enum from ._errors import DlpConfigurationError 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