dlp-io/dlp_io/_helper.py

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