mirror of
https://github.com/Strategic-Automation/violin.git
synced 2026-08-14 12:33:37 +02:00
50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
|
"""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)
|