Files

57 lines
2.4 KiB
Python

"""Transport-agnostic safeguards for an authorized reverse-shell PTY.
Wire ``send_line`` and ``read_until`` to an already established controller.
This template intentionally has no network listener and no file-transfer path:
use HTTP plus SHA-256 for files instead of base64 through a PTY.
"""
from __future__ import annotations
import hashlib
import secrets
from collections.abc import Callable
from pathlib import Path
MAX_PTY_LINE_BYTES = 320
class PtyShellController:
def __init__(self, send_line: Callable[[str], None], read_until: Callable[[str], str]) -> None:
self._send_line = send_line
self._read_until = read_until
def run(self, command: str) -> tuple[int, str]:
"""Send one short command and return labelled output plus its exit code."""
if not command.strip():
raise ValueError("command must not be empty")
if "\n" in command or "\r" in command:
raise ValueError("send one command line at a time")
if len(command.encode("utf-8")) > MAX_PTY_LINE_BYTES:
raise ValueError("PTY command is too long; deliver files over HTTP, not this channel")
marker = f"__VIOLIN_DONE_{secrets.token_hex(8)}__"
wrapped = f"{command}; printf '\\n{marker}=%s\\n' \"$?\""
if len(wrapped.encode("utf-8")) > MAX_PTY_LINE_BYTES:
raise ValueError("wrapped PTY command is too long; use an HTTP-delivered script")
self._send_line(wrapped)
transcript = self._read_until(marker)
marker_line = next(
(line for line in transcript.splitlines() if line.startswith(marker + "=")), ""
)
if not marker_line:
raise RuntimeError("controller did not receive the command completion marker")
try:
exit_code = int(marker_line.partition("=")[2])
except ValueError as exc:
raise RuntimeError("invalid command completion marker") from exc
output = "\n".join(
f"[VICTIM] {line}" for line in transcript.splitlines() if line != marker_line
)
return exit_code, output
def sha256_matches(local_path: str, victim_digest: str) -> bool:
"""Compare a local artifact with a SHA-256 value collected from the victim."""
with Path(local_path).open("rb") as artifact:
local_digest = hashlib.file_digest(artifact, "sha256").hexdigest()
return local_digest.lower() == victim_digest.strip().lower()