diff --git a/dlp_io.py b/dlp_io.py index d33aa4c..4e228cd 100644 --- a/dlp_io.py +++ b/dlp_io.py @@ -4,20 +4,24 @@ from __future__ import annotations import atexit import builtins +import errno +import hashlib import io import json import operator import os import queue +import re import shutil import subprocess +import tempfile import threading from dataclasses import dataclass from enum import Enum from multiprocessing.context import AuthenticationError from multiprocessing.connection import Client -__version__ = "0.1.1" +__version__ = "0.2.0" __all__ = [ "DlpConfig", @@ -28,9 +32,12 @@ __all__ = [ "DlpSession", "DlpSessionBusyError", "DlpTransportError", + "DlpWriteError", + "WriteChannel", "configure", "install_open_patch", "install_reader_patch", + "is_encrypted", "open", "shutdown", "uninstall_open_patch", @@ -65,6 +72,10 @@ class DlpSessionBusyError(DlpIoError): """A session already owns an active file stream.""" +class DlpWriteError(DlpIoError): + """No available write channel can produce an unencrypted file.""" + + # --------------------------------------------------------------------------- # Configuration @@ -797,7 +808,7 @@ def open( 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) + return _open_write(file, mode, buffering, encoding, errors, newline, closefd) def shutdown() -> None: @@ -865,7 +876,7 @@ def _reader_wrapper(open_read_func): newline=newline, line_buffering=buffering == 1, ) - return _ORIGINAL_OPEN(file, mode, buffering, encoding, errors, newline, closefd) + return _open_write(file, mode, buffering, encoding, errors, newline, closefd) return patched_open @@ -901,3 +912,282 @@ def uninstall_open_patch() -> None: builtins.open = _ORIGINAL_OPEN io.open = _ORIGINAL_IO_OPEN _installed = False + + +# --------------------------------------------------------------------------- +# Encryption detection and unencrypted writes +# +# DLP transparent encryption stores ciphertext on disk. Approved processes +# (the whitelisted Python) see plaintext through the filter driver, while +# non-whitelisted tools such as certutil see the raw on-disk bytes. File +# attributes (attrib, ADS, size) are identical between encrypted and plain +# files, so content comparison is the only reliable discriminator. + +_DETECT_HEAD_SIZE = 64 * 1024 +_RAW_HASH_PATTERN = re.compile(rb"^[ \t]*([0-9a-fA-F]{64})[ \t\r]*$", re.MULTILINE) + + +def _raw_sha256(path) -> bytes: + """SHA-256 of the raw on-disk bytes, read by non-whitelisted certutil. + + certutil must be spawned through ``cmd /c``: DLP hosts block approved + processes from starting certutil directly, and list-form arguments keep + non-ASCII paths intact through the cmd command line. + """ + try: + completed = subprocess.run( + ["cmd", "/c", "certutil", "-hashfile", path, "SHA256"], + capture_output=True, + **_helper_process_options(), + ) + except (OSError, subprocess.SubprocessError) as exc: + raise DlpIoError("raw certutil hash unavailable: %s" % exc) from exc + if completed.returncode != 0: + raise DlpIoError( + "certutil -hashfile failed with exit code %d" % completed.returncode + ) + match = _RAW_HASH_PATTERN.search(completed.stdout) + if match is None: + raise DlpIoError("certutil output did not contain a SHA-256 digest") + return bytes.fromhex(match.group(1).decode("ascii")) + + +def _native_sha256(path) -> bytes: + digest = hashlib.sha256() + with _ORIGINAL_OPEN(path, "rb") as stream: + for chunk in iter(lambda: stream.read(CHUNK_SIZE), b""): + digest.update(chunk) + return digest.digest() + + +def _native_head(path, size) -> bytes: + with _ORIGINAL_OPEN(path, "rb") as stream: + return stream.read(size) + + +def _helper_head(path, size) -> bytes: + raw = _get_default_session().open_raw(path) + try: + return raw.read(size) + finally: + raw.close() + + +def is_encrypted(path) -> bool: + """Return True when the file on disk holds DLP ciphertext. + + The raw on-disk digest (certutil) is compared with what the current + process sees. A mismatch means the file is encrypted and the current + process is approved. When both match the result is ambiguous — the file + is either plain, or encrypted while the current process is not approved + (packaged .exe) — so the first bytes seen by the approved helper are + compared with the first bytes seen by the current process. + """ + if os.name != "nt": + return False + normalized = os.fsdecode(os.fspath(path)) + native_hash = _native_sha256(normalized) + raw_hash = _raw_sha256(normalized) + if raw_hash != native_hash: + return True + return _helper_head(normalized, _DETECT_HEAD_SIZE) != _native_head( + normalized, _DETECT_HEAD_SIZE + ) + + +class WriteChannel(str, Enum): + """How path-based writes are routed to keep files unencrypted.""" + + DIRECT = "direct" + POWERSHELL = "powershell" + + +_WRITE_CHANNEL_ENV = "DLP_IO_WRITE_CHANNEL" +_WRITE_PROBE_PAYLOAD = b"dlp-io-write-probe\x00\x1a\xff" + bytes(range(256)) * 2 +_write_channel_lock = threading.RLock() +_write_channel_cache = {} + + +def _powershell_script(path, file_mode) -> str: + escaped = path.replace("'", "''") + return ( + "$ErrorActionPreference='Stop';" + "$input_stream=[Console]::OpenStandardInput();" + "$file_stream=[IO.File]::Open('" + escaped + "',[IO.FileMode]::" + file_mode + ");" + "$input_stream.CopyTo($file_stream);" + "$file_stream.Close()" + ) + + +class _PowerShellRelayRaw(io.RawIOBase): + """Binary stream written to disk by a non-whitelisted PowerShell child. + + Bytes travel through the stdin pipe (never through a DLP-encrypted temp + file), and the child's file writes are not encrypted because the child + is outside the DLP write-encryption policy. + """ + + def __init__(self, path, file_mode) -> None: + super().__init__() + try: + self._proc = subprocess.Popen( + ["powershell", "-NoProfile", "-Command", _powershell_script(path, file_mode)], + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + **_helper_process_options(), + ) + except OSError as exc: + raise DlpWriteError("failed to start powershell write relay: %s" % exc) from exc + + def writable(self) -> bool: + return True + + def write(self, data): + if self.closed: + raise ValueError("I/O operation on closed file") + try: + self._proc.stdin.write(data) + except (BrokenPipeError, OSError) as exc: + raise DlpWriteError("powershell write relay rejected data: %s" % exc) from exc + return len(data) + + def flush(self) -> None: + if self.closed: + raise ValueError("I/O operation on closed file") + stdin = self._proc.stdin + if stdin is None or stdin.closed: + return + try: + stdin.flush() + except (BrokenPipeError, OSError) as exc: + raise DlpWriteError("powershell write relay rejected data: %s" % exc) from exc + + def close(self) -> None: + if self.closed: + return + try: + try: + self._proc.stdin.close() + except (BrokenPipeError, OSError): + pass + returncode = self._proc.wait() + if returncode != 0: + detail = self._proc.stderr.read().decode("utf-8", "replace").strip() + raise DlpWriteError( + "powershell write relay failed with exit code %d: %s" + % (returncode, detail) + ) + finally: + super().close() + + +def _write_probe_path(suffix) -> str: + descriptor, probe = tempfile.mkstemp(prefix="dlp-io-probe-", suffix=suffix or None) + os.close(descriptor) + return probe + + +def _probe_direct_write(suffix) -> bool: + probe = _write_probe_path(suffix) + try: + with _ORIGINAL_OPEN(probe, "wb") as stream: + stream.write(_WRITE_PROBE_PAYLOAD) + try: + raw_hash = _raw_sha256(probe) + except DlpIoError: + # Without a raw channel there is no DLP filter to bypass. + return True + return raw_hash == hashlib.sha256(_WRITE_PROBE_PAYLOAD).digest() + finally: + os.unlink(probe) + + +def _probe_powershell_write(suffix) -> bool: + probe = _write_probe_path(suffix) + try: + try: + with _PowerShellRelayRaw(probe, "Create") as stream: + stream.write(_WRITE_PROBE_PAYLOAD) + return _raw_sha256(probe) == hashlib.sha256(_WRITE_PROBE_PAYLOAD).digest() + except (DlpIoError, OSError): + return False + finally: + os.unlink(probe) + + +def _select_write_channel(suffix) -> WriteChannel: + """Pick a write channel that lands unencrypted bytes on disk. + + Results are cached per file suffix because DLP write policies are + typically scoped by document type. ``DLP_IO_WRITE_CHANNEL`` forces a + channel and skips probing. + """ + override = os.environ.get(_WRITE_CHANNEL_ENV) + if override: + try: + return WriteChannel(override.strip().lower()) + except ValueError as exc: + valid = ", ".join(channel.value for channel in WriteChannel) + raise DlpConfigurationError( + "%s must be one of: %s" % (_WRITE_CHANNEL_ENV, valid) + ) from exc + key = suffix.lower() + with _write_channel_lock: + cached = _write_channel_cache.get(key) + if cached is not None: + return cached + if _probe_direct_write(suffix): + channel = WriteChannel.DIRECT + elif _probe_powershell_write(suffix): + channel = WriteChannel.POWERSHELL + else: + raise DlpWriteError( + "no write channel can keep %r files unencrypted on this host" + % (suffix or "") + ) + _write_channel_cache[key] = channel + return channel + + +def _wrap_writer(raw, mode, buffering, encoding, errors, newline): + binary = "b" in mode + if buffering == 0: + if not binary: + raw.close() + raise ValueError("can't have unbuffered text I/O") + return raw + 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.BufferedWriter(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_write(file, mode, buffering, encoding, errors, newline, closefd): + """Open a path for writing so the file stays unencrypted on disk.""" + path = os.fsdecode(os.fspath(file)) + if "x" in mode and os.path.exists(path): + raise FileExistsError(errno.EEXIST, "File exists", path) + channel = _select_write_channel(os.path.splitext(path)[1]) + if channel is WriteChannel.DIRECT: + return _ORIGINAL_OPEN(file, mode, buffering, encoding, errors, newline, closefd) + if not closefd: + raise ValueError("Cannot use closefd=False with file name") + if "w" in mode: + file_mode = "Create" + elif "a" in mode: + file_mode = "Append" + else: + file_mode = "CreateNew" + raw = _PowerShellRelayRaw(path, file_mode) + return _wrap_writer(raw, mode, buffering, encoding, errors, newline) diff --git a/docs/DLP_IO_LIBRARY.md b/docs/DLP_IO_LIBRARY.md index ed8ff85..adb506c 100644 --- a/docs/DLP_IO_LIBRARY.md +++ b/docs/DLP_IO_LIBRARY.md @@ -1,8 +1,8 @@ # dlp-io -`dlp-io` 为 Windows DLP 透明加密环境提供接近 `io.open` 的 Python 文件 API。非白名单进程负责业务逻辑和写入;文件读取由已获 DLP 白名单授权的 Python helper 完成,再通过经过认证的 Windows Named Pipe 返回明文字节。 +`dlp-io` 为 Windows DLP 透明加密环境提供接近 `io.open` 的 Python 文件 API。非白名单进程负责业务逻辑;文件读取由已获 DLP 白名单授权的 Python helper 完成,再通过经过认证的 Windows Named Pipe 返回明文字节;文件写入通过通道探测保证落盘为未加密格式。另提供 `is_encrypted()`,在读取前判断文件是否处于 DLP 加密状态。 -当前版本为 `dlp-io==0.1.1`,distribution 名是 `dlp-io`,import 名是 `dlp_io`。 +当前版本为 `dlp-io==0.2.0`,distribution 名是 `dlp-io`,import 名是 `dlp_io`。 ## 安装 @@ -11,7 +11,7 @@ registry 已开放匿名下载,一条命令即可安装;pip 会自动选择最匹配的 wheel:Windows 上 CPython 3.10–3.14 会得到对应的 pyd wheel(整个库编译为单个 `.pyd`),其余环境得到纯 Python wheel。 ```powershell -py -m pip install --index-url https://gitea.docker.antior.cn/api/packages/antior/pypi/simple dlp-io==0.1.1 +py -m pip install --index-url https://gitea.docker.antior.cn/api/packages/antior/pypi/simple dlp-io==0.2.0 ``` Package 页面: @@ -78,14 +78,50 @@ dlp_io.shutdown() 模式路由规则: -- 路径型 `r`、`rt`、`rb`:通过 helper 读取。 -- 路径型 `w`、`a`、`x`:使用原生 `open` 写入。 +- 路径型 `r`、`rt`、`rb`:通过 helper 读取,始终返回明文。 +- 路径型 `w`、`a`、`x`:经过写入通道路由(见下文「未加密写入」),保证落盘为未加密文件。 - 整数文件描述符或自定义 `opener`:使用原生 `open`。 - 路径型 `+` 更新模式:显式抛出 `io.UnsupportedOperation`。 - 路径读取配合 `closefd=False`:与标准 API 一样抛出 `ValueError`。 - `rb` 配合 `buffering=0`:返回 raw stream;默认返回 buffered reader。 - 文本模式的默认 encoding 与当前 Python `io.open` 一致;跨机器文件请显式传 `encoding="utf-8"`。 +## 加密状态判断 + +`is_encrypted(path)` 在读取前判断文件是否处于 DLP 加密状态,返回 `bool`: + +```python +import dlp_io + +if dlp_io.is_encrypted(r"D:\Protected\input.pptx"): + # 已加密:当前进程若是打包的 EXE(非白名单),必须经 helper 中转读取 + with dlp_io.open(r"D:\Protected\input.pptx", "rb") as stream: + data = stream.read() +else: + # 未加密:可用任意方式直接读取 + with dlp_io.open(r"D:\Protected\input.pptx", "rb") as stream: + data = stream.read() +``` + +判断原理:DLP 透明加密在磁盘上存密文,白名单进程读到明文、非白名单工具(certutil)读到原始盘内字节;加密文件与未加密文件的 attrib 属性、ADS、大小完全一致,只有内容可区分。`is_encrypted()` 先比较 certutil 原始哈希与当前进程视图哈希:不一致则文件已加密且当前进程是白名单;一致时存在歧义(文件未加密,或文件已加密但当前进程是非白名单 EXE),再比较 helper 视图与当前进程视图的前 64 KiB 定案。两种运行环境(Python/pyd 与打包 EXE)结果都正确。 + +代价:未加密文件需要一次完整 certutil 哈希、一次本地完整读取和一次 helper 前 64 KiB 读取;大文件请缓存判断结果,不要每次读取前重复调用。helper 未启动时会复用懒加载的默认 session,结束后照常 `dlp_io.shutdown()`。非 Windows 平台直接返回 `False`。 + +## 未加密写入 + +写入目标始终是「磁盘上保存未加密文件」: + +- 打包为 EXE(非白名单进程):原生直写天然落盘明文,选中直写通道。 +- Python 调试 / pyd(白名单进程):部分 DLP 策略会加密白名单进程的写入。首次写某个后缀时库会用临时探针文件实测直写是否落盘明文(用 certutil 原始哈希校验);若直写会被加密,自动改走 PowerShell 中转通道——字节经 stdin 管道传给非白名单的 powershell.exe 子进程写盘,不经过任何磁盘临时文件。通道按文件后缀缓存;两个通道都不可用时抛出 `DlpWriteError`,绝不静默写出密文。 + +可用环境变量跳过探测、强制指定通道(值为 `direct` 或 `powershell`): + +```powershell +$env:DLP_IO_WRITE_CHANNEL = "powershell" +``` + +探测按后缀缓存,注意 DLP 策略若按目录区分,探针结果可能与目标目录不同;此时建议用环境变量显式指定通道。 + ## 配置和生命周期 默认 session 在第一次读取时懒启动,并被后续读取复用: @@ -154,6 +190,7 @@ finally: - `DlpProtocolError`:协议版本、控制消息或响应结构无效。 - `DlpTransportError`:helper 在显式 EOF 前断开,或文件传输被截断。 - `DlpSessionBusyError`:同一个 session 已有活动 stream。 +- `DlpWriteError`:没有任何写入通道能落盘未加密文件,或 PowerShell 中转写入失败。 helper 启动或读取失败时,库会显式报错并且不回退到当前 EXE 直接读取。DLP 直读可能返回合法长度的密文;静默回退会把数据损坏伪装成成功。 diff --git a/tests/DLP加密文件.pptx b/tests/DLP加密文件.pptx new file mode 100644 index 0000000..ee7aed4 Binary files /dev/null and b/tests/DLP加密文件.pptx differ diff --git a/tests/DLP解密文件.pptx b/tests/DLP解密文件.pptx new file mode 100644 index 0000000..ee7aed4 Binary files /dev/null and b/tests/DLP解密文件.pptx differ diff --git a/tests/test_dlp_io_encryption.py b/tests/test_dlp_io_encryption.py new file mode 100644 index 0000000..f5eb725 --- /dev/null +++ b/tests/test_dlp_io_encryption.py @@ -0,0 +1,262 @@ +from __future__ import annotations + +import hashlib +import os +from pathlib import Path + +import pytest + +import dlp_io + + +pytestmark = pytest.mark.skipif(os.name != "nt", reason="DLP detection requires Windows") + +SAMPLES_DIR = Path(__file__).resolve().parent +ENCRYPTED_SAMPLE = SAMPLES_DIR / "DLP加密文件.pptx" +DECRYPTED_SAMPLE = SAMPLES_DIR / "DLP解密文件.pptx" + + +@pytest.fixture(autouse=True) +def reset_module_state(monkeypatch): + monkeypatch.delenv(dlp_io._WRITE_CHANNEL_ENV, raising=False) + monkeypatch.setattr(dlp_io, "_config", dlp_io.DlpConfig()) + dlp_io._write_channel_cache.clear() + dlp_io.shutdown() + yield + dlp_io._write_channel_cache.clear() + dlp_io.shutdown() + + +def _native_sha256(path: Path) -> bytes: + return hashlib.sha256(path.read_bytes()).digest() + + +def _dlp_host() -> bool: + """本机是否启用 DLP:加密样本的原始盘内字节与本地视图是否不一致。""" + return dlp_io._raw_sha256(str(ENCRYPTED_SAMPLE)) != _native_sha256(ENCRYPTED_SAMPLE) + + +# --------------------------------------------------------------------------- +# is_encrypted against the real DLP samples + + +def test_encrypted_sample_matches_host_dlp_capability() -> None: + assert dlp_io.is_encrypted(ENCRYPTED_SAMPLE) == _dlp_host() + + +def test_decrypted_sample_is_never_reported_encrypted() -> None: + assert dlp_io.is_encrypted(DECRYPTED_SAMPLE) is False + + +def test_is_encrypted_missing_file_raises() -> None: + with pytest.raises(FileNotFoundError): + dlp_io.is_encrypted(SAMPLES_DIR / "missing-sample.bin") + + +# --------------------------------------------------------------------------- +# is_encrypted decision branches (mocked channels) + + +def test_is_encrypted_short_circuits_when_raw_digest_differs(monkeypatch) -> None: + monkeypatch.setattr(dlp_io, "_native_sha256", lambda path: b"\x01" * 32) + monkeypatch.setattr(dlp_io, "_raw_sha256", lambda path: b"\x02" * 32) + monkeypatch.setattr( + dlp_io, + "_get_default_session", + lambda: pytest.fail("helper must not start when digests differ"), + ) + + assert dlp_io.is_encrypted("payload.bin") is True + + +@pytest.mark.parametrize( + ("helper_head", "native_head", "expected"), + [ + (b"\x50\x4b\x03\x04", b"\x18\x1b\x03\x1a", True), + (b"\x50\x4b\x03\x04", b"\x50\x4b\x03\x04", False), + ], +) +def test_is_encrypted_ambiguous_case_compares_helper_head( + monkeypatch, helper_head: bytes, native_head: bytes, expected: bool +) -> None: + digest = b"\x01" * 32 + monkeypatch.setattr(dlp_io, "_native_sha256", lambda path: digest) + monkeypatch.setattr(dlp_io, "_raw_sha256", lambda path: digest) + monkeypatch.setattr(dlp_io, "_native_head", lambda path, size: native_head) + monkeypatch.setattr(dlp_io, "_helper_head", lambda path, size: helper_head) + + assert dlp_io.is_encrypted("payload.bin") is expected + + +# --------------------------------------------------------------------------- +# Write channel selection + + +def test_select_write_channel_honours_probes(monkeypatch, tmp_path) -> None: + monkeypatch.setattr(dlp_io, "_probe_direct_write", lambda suffix: True) + monkeypatch.setattr( + dlp_io, + "_probe_powershell_write", + lambda suffix: pytest.fail("powershell probe must be skipped"), + ) + assert dlp_io._select_write_channel(".bin") is dlp_io.WriteChannel.DIRECT + + monkeypatch.setattr(dlp_io, "_probe_direct_write", lambda suffix: False) + monkeypatch.setattr(dlp_io, "_probe_powershell_write", lambda suffix: True) + assert dlp_io._select_write_channel(".pptx") is dlp_io.WriteChannel.POWERSHELL + + +def test_select_write_channel_raises_when_no_channel_works(monkeypatch) -> None: + monkeypatch.setattr(dlp_io, "_probe_direct_write", lambda suffix: False) + monkeypatch.setattr(dlp_io, "_probe_powershell_write", lambda suffix: False) + + with pytest.raises(dlp_io.DlpWriteError): + dlp_io._select_write_channel(".bin") + + +def test_select_write_channel_caches_per_suffix(monkeypatch) -> None: + calls = [] + monkeypatch.setattr( + dlp_io, "_probe_direct_write", lambda suffix: calls.append(suffix) or True + ) + + assert dlp_io._select_write_channel(".BIN") is dlp_io.WriteChannel.DIRECT + assert dlp_io._select_write_channel(".bin") is dlp_io.WriteChannel.DIRECT + assert len(calls) == 1 + + +def test_write_channel_env_override(monkeypatch) -> None: + monkeypatch.setenv(dlp_io._WRITE_CHANNEL_ENV, "PowerShell") + monkeypatch.setattr( + dlp_io, + "_probe_direct_write", + lambda suffix: pytest.fail("override must skip probing"), + ) + assert dlp_io._select_write_channel(".bin") is dlp_io.WriteChannel.POWERSHELL + + monkeypatch.setenv(dlp_io._WRITE_CHANNEL_ENV, "bogus") + with pytest.raises(dlp_io.DlpConfigurationError): + dlp_io._select_write_channel(".bin") + + +def test_real_probe_selects_a_working_channel(tmp_path) -> None: + try: + channel = dlp_io._select_write_channel(".bin") + except dlp_io.DlpWriteError: + pytest.skip("host cannot produce unencrypted writes") + + direct_ok = dlp_io._probe_direct_write(".bin") + expected = dlp_io.WriteChannel.DIRECT if direct_ok else dlp_io.WriteChannel.POWERSHELL + assert channel is expected + + +# --------------------------------------------------------------------------- +# PowerShell relay writer + + +def test_powershell_relay_writes_exact_binary_and_lands_plaintext(tmp_path) -> None: + payload = bytes(range(256)) * 64 + b"\x00\x1a\xff" + target = tmp_path / "relay输出.bin" + + with dlp_io._PowerShellRelayRaw(str(target), "Create") as raw: + assert raw.writable() + assert raw.write(payload) == len(payload) + + assert target.read_bytes() == payload + assert dlp_io._raw_sha256(str(target)) == hashlib.sha256(payload).digest() + + +def test_powershell_relay_append_preserves_existing_content(tmp_path) -> None: + target = tmp_path / "relay.bin" + + with dlp_io._PowerShellRelayRaw(str(target), "Create") as raw: + raw.write(b"ab") + with dlp_io._PowerShellRelayRaw(str(target), "Append") as raw: + raw.write(b"cd") + + assert target.read_bytes() == b"abcd" + + +def test_powershell_relay_create_new_fails_on_existing_file(tmp_path) -> None: + target = tmp_path / "relay.bin" + target.write_bytes(b"existing") + + with pytest.raises(dlp_io.DlpWriteError): + with dlp_io._PowerShellRelayRaw(str(target), "CreateNew"): + pass + + +def test_open_write_routes_through_relay_when_selected(monkeypatch, tmp_path) -> None: + monkeypatch.setattr( + dlp_io, "_select_write_channel", lambda suffix: dlp_io.WriteChannel.POWERSHELL + ) + payload = bytes(range(256)) * 32 + target = tmp_path / "routed.bin" + + with dlp_io.open(target, "wb") as file: + file.write(payload) + + assert target.read_bytes() == payload + + +def test_open_text_write_routes_through_relay(monkeypatch, tmp_path) -> None: + monkeypatch.setattr( + dlp_io, "_select_write_channel", lambda suffix: dlp_io.WriteChannel.POWERSHELL + ) + payload = "中文内容\n第二行".encode("utf-8") + target = tmp_path / "routed.txt" + + with dlp_io.open(target, "w", encoding="utf-8", newline="") as file: + file.write(payload.decode("utf-8")) + with dlp_io.open(target, "a", encoding="utf-8", newline="") as file: + file.write("追加") + + assert target.read_bytes() == payload + "追加".encode("utf-8") + + +def test_open_exclusive_create_is_atomic_before_channel_selection( + monkeypatch, tmp_path +) -> None: + monkeypatch.setattr( + dlp_io, + "_select_write_channel", + lambda suffix: pytest.fail("existing target must fail before probing"), + ) + target = tmp_path / "existing.bin" + target.write_bytes(b"data") + + with pytest.raises(FileExistsError): + dlp_io.open(target, "xb") + + assert target.read_bytes() == b"data" + + +def test_open_write_direct_channel_uses_native_open(monkeypatch, tmp_path) -> None: + monkeypatch.setattr( + dlp_io, "_select_write_channel", lambda suffix: dlp_io.WriteChannel.DIRECT + ) + target = tmp_path / "direct.bin" + + with dlp_io.open(target, "wb") as file: + assert not isinstance(file, dlp_io._PowerShellRelayRaw) + file.write(b"native") + + assert target.read_bytes() == b"native" + + +def test_patched_open_write_uses_same_channel_routing(monkeypatch, tmp_path) -> None: + import builtins + + monkeypatch.setattr( + dlp_io, "_select_write_channel", lambda suffix: dlp_io.WriteChannel.POWERSHELL + ) + target = tmp_path / "patched.bin" + + dlp_io.install_open_patch() + try: + with builtins.open(target, "wb") as file: + file.write(b"patched-write") + finally: + dlp_io.uninstall_open_patch() + + assert target.read_bytes() == b"patched-write"