dlp-io/tests/test_dlp_io_encryption.py

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"