"""Safety contract for the documented PTY shell controller template."""

from __future__ import annotations

import hashlib
import importlib.util
from pathlib import Path

import pytest

TEMPLATE = Path(__file__).parents[2] / "skills" / "pentest" / "templates" / "shell_ctrl.py"
SPEC = importlib.util.spec_from_file_location("violin_shell_ctrl_template", TEMPLATE)
assert SPEC and SPEC.loader
shell_ctrl = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(shell_ctrl)


def test_controller_labels_victim_output_and_parses_exit(monkeypatch):
    sent: list[str] = []
    monkeypatch.setattr(shell_ctrl.secrets, "token_hex", lambda _: "fixed")
    controller = shell_ctrl.PtyShellController(
        sent.append,
        lambda marker: f"uid=1000\n{marker}=7\n",
    )

    exit_code, output = controller.run("id")

    assert exit_code == 7
    assert output == "[VICTIM] uid=1000"
    assert sent == ["id; printf '\\n__VIOLIN_DONE_fixed__=%s\\n' \"$?\""]


def test_controller_refuses_long_or_multiline_pty_delivery():
    controller = shell_ctrl.PtyShellController(lambda _: None, lambda _: "")

    with pytest.raises(ValueError, match="too long"):
        controller.run("x" * shell_ctrl.MAX_PTY_LINE_BYTES)
    with pytest.raises(ValueError, match="one command line"):
        controller.run("id\nuname -a")


def test_sha256_matches_requires_exact_digest(tmp_path):
    artifact = tmp_path / "payload.py"
    content = b"print('safe')\n"
    artifact.write_bytes(content)
    digest = hashlib.sha256(content).hexdigest()

    assert shell_ctrl.sha256_matches(str(artifact), digest)
    assert not shell_ctrl.sha256_matches(str(artifact), "0" * 64)