cyber-shell-wrapper 0.1.0__py3-none-any.whl

cyber_shell/__init__.py

@@ -0,0 +1,5 @@
+"""Cyber Shell package."""
+
+__all__ = ["__version__"]
+
+__version__ = "0.1.0"

cyber_shell/assembler.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import re
+import shlex
+from dataclasses import dataclass
+
+from .config import AppConfig
+from .models import ShellEvent
+
+
+INTERACTIVE_COMMANDS = {
+    "alsamixer",
+    "ftp",
+    "htop",
+    "less",
+    "man",
+    "more",
+    "mongo",
+    "mysql",
+    "nano",
+    "nmtui",
+    "psql",
+    "python",
+    "python3",
+    "redis-cli",
+    "sftp",
+    "ssh",
+    "sqlite3",
+    "telnet",
+    "tig",
+    "tmux",
+    "top",
+    "vi",
+    "view",
+    "vim",
+    "watch",
+}
+
+PREFIX_WRAPPERS = {
+    "builtin",
+    "chronic",
+    "command",
+    "env",
+    "exec",
+    "nohup",
+    "stdbuf",
+    "time",
+}
+
+ANSI_ESCAPE_RE = re.compile(
+    r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~]|\][^\x1b\x07]*(?:\x07|\x1b\\))"
+)
+
+
+@dataclass(slots=True)
+class ActiveCommand:
+    started_at: str
+    cmd: str
+    output_buffer: bytearray
+    truncated: bool = False
+
+
+class EventAssembler:
+    def __init__(self, config: AppConfig, session_id: str) -> None:
+        self._config = config
+        self._session_id = session_id
+        self._seq = 0
+        self._current: ActiveCommand | None = None
+
+    def start_command(self, started_at: str, cmd: str) -> None:
+        self._current = ActiveCommand(
+            started_at=started_at,
+            cmd=cmd.strip(),
+            output_buffer=bytearray(),
+        )
+
+    def append_output(self, chunk: bytes) -> None:
+        if self._current is None or not chunk:
+            return
+        remaining = self._config.max_output_bytes - len(self._current.output_buffer)
+        if remaining <= 0:
+            self._current.truncated = True
+            return
+        self._current.output_buffer.extend(chunk[:remaining])
+        if len(chunk) > remaining:
+            self._current.truncated = True
+
+    def finish_command(
+        self,
+        *,
+        finished_at: str,
+        exit_code: int,
+        cwd: str,
+    ) -> ShellEvent | None:
+        current = self._current
+        self._current = None
+        if current is None or not current.cmd:
+            return None
+
+        self._seq += 1
+        return ShellEvent(
+            session_id=self._session_id,
+            hostname=self._config.hostname,
+            shell="bash",
+            seq=self._seq,
+            cwd=cwd,
+            cmd=current.cmd,
+            exit_code=exit_code,
+            output=_sanitize_output(
+                current.output_buffer.decode("utf-8", errors="replace")
+            ),
+            output_truncated=current.truncated,
+            started_at=current.started_at,
+            finished_at=finished_at,
+            is_interactive=is_interactive_command(current.cmd),
+            metadata=dict(self._config.metadata),
+        )
+
+
+def is_interactive_command(command: str) -> bool:
+    executable = _extract_command_name(command)
+    if executable is None:
+        return False
+    return executable in INTERACTIVE_COMMANDS
+
+
+def _extract_command_name(command: str) -> str | None:
+    try:
+        tokens = shlex.split(command, posix=True)
+    except ValueError:
+        return None
+
+    index = 0
+    while index < len(tokens):
+        token = tokens[index]
+        if _looks_like_env_assignment(token):
+            index += 1
+            continue
+        if token == "sudo":
+            index += 1
+            while index < len(tokens):
+                sudo_token = tokens[index]
+                if sudo_token == "--":
+                    index += 1
+                    break
+                if not sudo_token.startswith("-"):
+                    break
+                index += 1
+                if sudo_token in {"-g", "-h", "-p", "-u"} and index < len(tokens):
+                    index += 1
+            continue
+        if token == "env":
+            index += 1
+            while index < len(tokens) and _looks_like_env_assignment(tokens[index]):
+                index += 1
+            continue
+        if token in PREFIX_WRAPPERS:
+            index += 1
+            continue
+        return token
+    return None
+
+
+def _looks_like_env_assignment(token: str) -> bool:
+    if "=" not in token or token.startswith("="):
+        return False
+    name, _ = token.split("=", 1)
+    return name.replace("_", "A").isalnum()
+
+
+def _sanitize_output(value: str) -> str:
+    cleaned = ANSI_ESCAPE_RE.sub("", value)
+    cleaned = cleaned.replace("\r\n", "\n").replace("\r", "\n")
+    return cleaned

cyber_shell/cli.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import argparse
+import sys
+
+from .config import default_config_text, has_runtime_overrides, load_config, persist_config
+from .logging_utils import configure_logging
+from .mock_endpoint import run_mock_endpoint
+from .telemetry import TelemetryClient
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="cyber-shell",
+        description="Local PTY shell wrapper for cyber range telemetry.",
+    )
+    parser.add_argument(
+        "--config",
+        help="Global config path shortcut for default start command.",
+    )
+    parser.add_argument(
+        "--endpoint-url",
+        help="Override telemetry endpoint for the wrapped shell session.",
+    )
+    parser.add_argument(
+        "--api-key",
+        help="Override telemetry API key for the wrapped shell session.",
+    )
+
+    subparsers = parser.add_subparsers(dest="command")
+
+    start_parser = subparsers.add_parser("start", help="Start wrapped Bash session.")
+    start_parser.add_argument(
+        "--config",
+        help="Path to config.yaml. Defaults to ~/.config/cyber-shell/config.yaml",
+    )
+    start_parser.add_argument(
+        "--endpoint-url",
+        help="Override telemetry endpoint for this session.",
+    )
+    start_parser.add_argument(
+        "--api-key",
+        help="Override telemetry API key for this session.",
+    )
+
+    mock_parser = subparsers.add_parser(
+        "mock-endpoint",
+        help="Run a local mock telemetry endpoint.",
+    )
+    mock_parser.add_argument("--host", default="127.0.0.1")
+    mock_parser.add_argument("--port", type=int, default=8080)
+    mock_parser.add_argument("--api-key")
+
+    subparsers.add_parser(
+        "print-default-config",
+        help="Print a default config.yaml template.",
+    )
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    command = args.command or "start"
+
+    if command == "print-default-config":
+        sys.stdout.write(default_config_text())
+        return 0
+
+    if command == "mock-endpoint":
+        return run_mock_endpoint(args.host, args.port, args.api_key)
+
+    config_path = getattr(args, "config", None)
+    config = load_config(config_path)
+    if getattr(args, "endpoint_url", None):
+        config.endpoint_url = args.endpoint_url
+    if getattr(args, "api_key", None):
+        config.api_key = args.api_key
+    logger = configure_logging(config.state_dir)
+
+    if command == "start":
+        if has_runtime_overrides(
+            {
+                "endpoint_url": getattr(args, "endpoint_url", None),
+                "api_key": getattr(args, "api_key", None),
+            }
+        ):
+            persisted_path = persist_config(config)
+            print(
+                f"cyber-shell: updated config -> {persisted_path}",
+                file=sys.stderr,
+            )
+        if config.endpoint_url:
+            print(
+                f"cyber-shell: telemetry -> {config.endpoint_url}",
+                file=sys.stderr,
+            )
+        else:
+            print(
+                "cyber-shell: telemetry disabled; set endpoint_url in config, "
+                "export CYBER_SHELL_ENDPOINT_URL, or pass --endpoint-url",
+                file=sys.stderr,
+            )
+        if config.endpoint_url and not config.api_key:
+            print(
+                "cyber-shell: telemetry has no API key; pass --api-key or "
+                "export CYBER_SHELL_API_KEY if the endpoint requires auth",
+                file=sys.stderr,
+            )
+        from .shell_wrapper import ShellWrapper
+
+        telemetry = TelemetryClient(config, logger)
+        try:
+            wrapper = ShellWrapper(config, telemetry, logger)
+            return wrapper.run()
+        finally:
+            telemetry.close()
+
+    parser.print_help()
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

cyber_shell/config.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+import os
+import socket
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+DEFAULT_CONFIG_PATH = Path.home() / ".config" / "cyber-shell" / "config.yaml"
+DEFAULT_STATE_DIR = Path.home() / ".local" / "state" / "cyber-shell"
+PERSISTED_ENV_KEYS = {
+    "CYBER_SHELL_CONFIG",
+    "CYBER_SHELL_ENDPOINT_URL",
+    "CYBER_SHELL_API_KEY",
+    "CYBER_SHELL_TIMEOUT_MS",
+    "CYBER_SHELL_RETRY_MAX",
+    "CYBER_SHELL_RETRY_BACKOFF_MS",
+    "CYBER_SHELL_MAX_OUTPUT_BYTES",
+    "CYBER_SHELL_QUEUE_SIZE",
+    "CYBER_SHELL_SHELL_PATH",
+    "CYBER_SHELL_HOSTNAME",
+    "CYBER_SHELL_STATE_DIR",
+}
+
+
+@dataclass(slots=True)
+class AppConfig:
+    endpoint_url: str | None = None
+    api_key: str | None = None
+    timeout_ms: int = 3000
+    retry_max: int = 3
+    retry_backoff_ms: int = 1000
+    max_output_bytes: int = 262144
+    queue_size: int = 256
+    shell_path: str = "/bin/bash"
+    state_dir: Path = DEFAULT_STATE_DIR
+    config_path: Path = DEFAULT_CONFIG_PATH
+    hostname: str = field(default_factory=socket.gethostname)
+    metadata: dict[str, str] = field(default_factory=dict)
+
+    def ensure_state_dir(self) -> Path:
+        self.state_dir.mkdir(parents=True, exist_ok=True)
+        return self.state_dir
+
+
+def load_config(path: str | Path | None = None) -> AppConfig:
+    config_path = Path(
+        path
+        or os.environ.get("CYBER_SHELL_CONFIG")
+        or DEFAULT_CONFIG_PATH
+    ).expanduser()
+    data: dict[str, object] = {}
+    if config_path.exists():
+        data = _parse_simple_yaml(config_path.read_text(encoding="utf-8"))
+
+    config = AppConfig(
+        endpoint_url=_env_or_data("CYBER_SHELL_ENDPOINT_URL", data, "endpoint_url"),
+        api_key=_env_or_data("CYBER_SHELL_API_KEY", data, "api_key"),
+        timeout_ms=_coerce_int(
+            _env_or_data("CYBER_SHELL_TIMEOUT_MS", data, "timeout_ms"), 3000
+        ),
+        retry_max=_coerce_int(
+            _env_or_data("CYBER_SHELL_RETRY_MAX", data, "retry_max"), 3
+        ),
+        retry_backoff_ms=_coerce_int(
+            _env_or_data("CYBER_SHELL_RETRY_BACKOFF_MS", data, "retry_backoff_ms"),
+            1000,
+        ),
+        max_output_bytes=_coerce_int(
+            _env_or_data("CYBER_SHELL_MAX_OUTPUT_BYTES", data, "max_output_bytes"),
+            262144,
+        ),
+        queue_size=_coerce_int(
+            _env_or_data("CYBER_SHELL_QUEUE_SIZE", data, "queue_size"), 256
+        ),
+        shell_path=str(
+            _env_or_data("CYBER_SHELL_SHELL_PATH", data, "shell_path") or "/bin/bash"
+        ),
+        state_dir=Path(
+            _env_or_data("CYBER_SHELL_STATE_DIR", data, "state_dir")
+            or DEFAULT_STATE_DIR
+        ).expanduser(),
+        config_path=config_path,
+        hostname=str(
+            _env_or_data("CYBER_SHELL_HOSTNAME", data, "hostname")
+            or socket.gethostname()
+        ),
+        metadata=_coerce_metadata(data.get("metadata")),
+    )
+    config.ensure_state_dir()
+    return config
+
+
+def default_config_text() -> str:
+    return _serialize_config(
+        AppConfig(
+            endpoint_url="http://127.0.0.1:8080/api/terminal-events",
+            api_key="replace-me",
+            timeout_ms=3000,
+            retry_max=3,
+            retry_backoff_ms=1000,
+            max_output_bytes=262144,
+            queue_size=256,
+            shell_path="/bin/bash",
+            metadata={"hostname_group": "kali-lab"},
+        )
+    )
+
+
+def persist_config(config: AppConfig) -> Path:
+    config.config_path.parent.mkdir(parents=True, exist_ok=True)
+    config.config_path.write_text(_serialize_config(config), encoding="utf-8")
+    if os.name == "posix":
+        os.chmod(config.config_path, 0o600)
+    return config.config_path
+
+
+def has_runtime_overrides(cli_args: dict[str, object] | None = None) -> bool:
+    if cli_args:
+        if cli_args.get("endpoint_url") or cli_args.get("api_key"):
+            return True
+    return any(key in os.environ for key in PERSISTED_ENV_KEYS)
+
+
+def _env_or_data(env_name: str, data: dict[str, object], key: str) -> object:
+    if env_name in os.environ:
+        return os.environ[env_name]
+    return data.get(key)
+
+
+def _coerce_int(value: object, default: int) -> int:
+    if value in (None, ""):
+        return default
+    try:
+        return int(str(value))
+    except ValueError:
+        return default
+
+
+def _coerce_metadata(value: object) -> dict[str, str]:
+    if not isinstance(value, dict):
+        return {}
+    return {str(key): str(inner) for key, inner in value.items()}
+
+
+def _serialize_config(config: AppConfig) -> str:
+    lines = [
+        f'endpoint_url: {_yaml_string(config.endpoint_url)}'
+        if config.endpoint_url is not None
+        else "endpoint_url: null",
+        f'api_key: {_yaml_string(config.api_key)}'
+        if config.api_key is not None
+        else "api_key: null",
+        f"timeout_ms: {config.timeout_ms}",
+        f"retry_max: {config.retry_max}",
+        f"retry_backoff_ms: {config.retry_backoff_ms}",
+        f"max_output_bytes: {config.max_output_bytes}",
+        f"queue_size: {config.queue_size}",
+        f'shell_path: {_yaml_string(config.shell_path)}',
+    ]
+    if config.metadata:
+        lines.append("metadata:")
+        for key, value in config.metadata.items():
+            lines.append(f"  {key}: {_yaml_string(value)}")
+    return "\n".join(lines) + "\n"
+
+
+def _yaml_string(value: object) -> str:
+    escaped = str(value).replace("\\", "\\\\").replace('"', '\\"')
+    return f'"{escaped}"'
+
+
+def _parse_simple_yaml(text: str) -> dict[str, object]:
+    result: dict[str, object] = {}
+    stack: list[tuple[int, dict[str, object]]] = [(-1, result)]
+
+    for raw_line in text.splitlines():
+        if not raw_line.strip():
+            continue
+        stripped = raw_line.lstrip()
+        if stripped.startswith("#"):
+            continue
+
+        indent = len(raw_line) - len(stripped)
+        while len(stack) > 1 and indent <= stack[-1][0]:
+            stack.pop()
+
+        if ":" not in stripped:
+            continue
+        key, raw_value = stripped.split(":", 1)
+        key = key.strip()
+        value = raw_value.strip()
+
+        parent = stack[-1][1]
+        if not value:
+            child: dict[str, object] = {}
+            parent[key] = child
+            stack.append((indent, child))
+            continue
+
+        parent[key] = _parse_scalar(value)
+
+    return result
+
+
+def _parse_scalar(value: str) -> object:
+    if value.startswith(("'", '"')) and value.endswith(("'", '"')) and len(value) >= 2:
+        return value[1:-1]
+    lowered = value.lower()
+    if lowered == "true":
+        return True
+    if lowered == "false":
+        return False
+    if lowered in {"null", "none"}:
+        return None
+    try:
+        return int(value)
+    except ValueError:
+        return value

cyber_shell/logging_utils.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import logging
+from logging.handlers import RotatingFileHandler
+from pathlib import Path
+
+
+def configure_logging(state_dir: Path) -> logging.Logger:
+    logger = logging.getLogger("cyber-shell")
+    if logger.handlers:
+        return logger
+
+    logger.setLevel(logging.INFO)
+    handler = RotatingFileHandler(
+        state_dir / "cyber-shell.log",
+        maxBytes=1_000_000,
+        backupCount=3,
+        encoding="utf-8",
+    )
+    handler.setFormatter(
+        logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s")
+    )
+    logger.addHandler(handler)
+    logger.propagate = False
+    return logger