sandbox-cli 0.2.26__py3-none-any.whl

sandbox_cli/utils/compiler/ssh.py

@@ -0,0 +1,203 @@
+import os
+import random
+import shutil
+import string
+import sys
+from pathlib import Path, PurePosixPath
+from typing import Any, cast
+
+from asyncssh import (
+    HostKeyNotVerifiable,
+    SSHClientConnection,
+    SSHClientConnectionOptions,
+    SSHClientProcess,
+    SSHCompletedProcess,
+    SSHReader,
+    connect,
+)
+
+from sandbox_cli.console import console
+from sandbox_cli.internal.config import settings
+from sandbox_cli.utils.compiler.abc import AbstractCompiler
+
+
+class RemoteCompiler(AbstractCompiler):
+    host: str
+    username: str
+    password: str
+    client: SSHClientConnection | None = None
+    tmp_directory: PurePosixPath
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.host = settings.sandbox[0].host
+        self.username = settings.sandbox[0].ssh.username
+        self.password = settings.sandbox[0].ssh.password
+
+    def _generate_random_string(self) -> str:
+        return "".join(random.choice(string.ascii_lowercase) for _ in range(8))
+
+    async def _init_ssh_client(self) -> None:
+        if not self.client:
+            try:
+                self.client = await connect(
+                    host=self.host,
+                    username=self.username,
+                    password=self.username,
+                    options=SSHClientConnectionOptions(public_key_auth=False),
+                )
+            except HostKeyNotVerifiable:
+                console.error(f"Can't verify ssh-key. Execute 'ssh {self.username}@{self.host}' and type 'yes'")
+                sys.exit(1)
+
+    async def _run_command(self, command: str, is_sudo: bool = False) -> SSHCompletedProcess:
+        await self._init_ssh_client()
+        assert self.client is not None
+
+        if is_sudo:
+            command = f"echo {self.password} | sudo -k -S {command}"
+
+        return await self.client.run(command)
+
+    async def _run_stream_command(self, command: str, is_sudo: bool = False) -> SSHClientProcess[Any]:
+        await self._init_ssh_client()
+        assert self.client is not None
+
+        if is_sudo:
+            command = f"echo {self.password} | sudo -k -S {command}"
+
+        async with self.client.create_process(command) as process:
+            async for line in cast(SSHReader[bytes], process.stdout):
+                print(f"{line.decode() if isinstance(line, bytes) else line}", end="")
+
+        return process
+
+    async def _create_tmp_directory(self) -> None:
+        self.tmp_directory = PurePosixPath("/tmp/sandbox-cli") / self._generate_random_string()
+        await self._run_command(f"mkdir -p {self.tmp_directory}")
+
+    async def _upload_rules(self, rules_dir: Path) -> None:
+        assert self.client is not None
+
+        arcname = shutil.make_archive("rules", "zip", rules_dir)
+        async with self.client.start_sftp_client() as ftp:
+            await ftp.put(arcname, f"{self.tmp_directory}")
+        os.remove(arcname)
+
+        await self._run_command(
+            f"""
+                mkdir -p {self.tmp_directory / "rules"} &&
+                mkdir -p {self.tmp_directory / "compiled-rules"} &&
+                unzip -d {self.tmp_directory / "rules"} {self.tmp_directory / "rules.zip"} >/dev/null
+            """
+        )
+
+    async def _cleanup(self, rules_dir: Path) -> None:
+        await self._run_command(f"rm -rf {self.tmp_directory}")
+        shutil.rmtree(rules_dir / "taxonomy")
+
+    async def _compile_rules_on_server(self) -> None:
+        # generate random container name for avoid collision with several peoples
+        result = await self._run_command(
+            command=f'ctr run --rm --memory-limit=1000000000 \
+            --mount type=bind,src={self.tmp_directory / "compiled-rules"},dst=/compiled-rules,options=rbind:rw \
+            --mount type=bind,src={self.tmp_directory / "rules"},dst=/rules,options=rbind:rw \
+            "{settings.docker.path}:{settings.docker.image_tag}" {self._generate_random_string()} package-builder correlation:compile -r /rules -c /compiled-rules',
+            is_sudo=True,
+        )
+        if result.exit_status:
+            console.log("failed to compile rules", style="bold red")
+            console.print(result.stderr)
+            sys.exit(1)
+
+    async def _download_compiled_rules(self) -> bytes:
+        assert self.client is not None
+
+        await self._run_command(
+            f"tar -C {self.tmp_directory / 'compiled-rules'} -czvf {self.tmp_directory / 'compiled-rules.tar.gz'} event_normalization_graph.json event_correlation_graph.json"
+        )
+
+        async with self.client.start_sftp_client() as ftp:
+            fd = await ftp.open(path=self.tmp_directory / "compiled-rules.tar.gz", pflags_or_mode="rb")
+            data: bytes = await fd.read()
+
+        return data
+
+    async def pull_image(self) -> None:
+        """Update image on remote server"""
+
+        process = await self._run_command(
+            f'ctr image pull --user "{settings.docker.username}:{settings.docker.token}" "{settings.docker.path}:{settings.docker.image_tag}"',
+            is_sudo=True,
+        )
+
+        if process.exit_status != 0:
+            console.error("Failed to update docker image on server")
+
+            if process.stdout:
+                value = process.stdout
+                console.log(f"{value.decode() if isinstance(value, bytes) else value}")
+            if process.stderr:
+                value = process.stderr
+                console.log(f"{value.decode() if isinstance(value, bytes) else value}")
+
+            sys.exit(1)
+
+        console.info("Docker image successfully updated on server")
+
+    async def compile_rules(self, rules_dir: Path, compiled_rules_dir: Path | None) -> bytes | None:
+        """
+        Compile rules on remote server
+        :params rules_dir path with rules
+        """
+
+        rules_dir = rules_dir.expanduser().resolve()
+        if not rules_dir.is_dir():
+            console.error("Invalid rules directory: {rules_dir}")
+            sys.exit(1)
+
+        # always take new version of taxonomy
+        shutil.copytree(rules_dir.parent / "taxonomy", rules_dir / "taxonomy", dirs_exist_ok=True)
+
+        compiled_rules: bytes | None = None
+        try:
+            # prepare folder
+            await self._create_tmp_directory()
+
+            # upload rules for compilation
+            await self._upload_rules(rules_dir)
+
+            await self._compile_rules_on_server()
+            compiled_rules = await self._download_compiled_rules()
+        finally:
+            # cleanup
+            await self._cleanup(rules_dir)
+
+        return compiled_rules
+
+    async def test_rules(self, root_rules_dir: Path, container_rules_dir: Path) -> bool:
+        # always take new version of taxonomy
+        shutil.copytree(root_rules_dir.parent / "taxonomy", root_rules_dir / "taxonomy", dirs_exist_ok=True)
+
+        try:
+            # prepare folder
+            await self._create_tmp_directory()
+
+            # upload rules for compilation
+            await self._upload_rules(root_rules_dir)
+
+            await self._compile_rules_on_server()
+
+            # generate random container name for avoid collision with several peoples
+            process = await self._run_stream_command(
+                f'ctr run --rm --memory-limit=1000000000 \
+                --mount type=bind,src={self.tmp_directory / "compiled-rules"},dst=/compiled-rules,options=rbind:rw \
+                --mount type=bind,src={self.tmp_directory / "rules" / container_rules_dir},dst=/rules,options=rbind:rw \
+                "{settings.docker.path}:{settings.docker.image_tag}" {self._generate_random_string()} package-builder correlation:test -r /rules -c /compiled-rules',
+                is_sudo=True,
+            )
+        finally:
+            # cleanup
+            await self._cleanup(root_rules_dir)
+
+        return process.exit_status == 0

sandbox_cli/utils/downloader/__init__.py

@@ -0,0 +1,178 @@
+import asyncio
+from collections.abc import Coroutine
+from io import BytesIO
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
+import aiofiles
+import zstandard
+from ptsandbox import Sandbox
+from ptsandbox.models import (
+    ArtifactType,
+    LogType,
+    SandboxBaseTaskResponse,
+    SandboxFileNotFoundException,
+)
+from rich.markup import escape
+from rich.progress import Progress, TaskID
+
+from sandbox_cli.console import console
+
+semaphore = asyncio.Semaphore(value=12)
+
+
+async def _save_artifact(
+    scan_id: UUID,
+    sandbox: Sandbox,
+    path: Path,
+    file_uri: str,
+    decompress: bool = False,
+    progress: Progress | None = None,
+    idx: str | None = None,
+    image: str | None = None,
+    link: str | None = None,
+) -> None:
+    _, uri = file_uri.split(":")
+    if not uri:
+        return
+
+    # sanitize path
+    path = Path(str(path).replace(" ", "_"))
+    task_id: TaskID = None  # type: ignore
+
+    async with semaphore:
+        if progress:
+            if idx and image and link:
+                task_id = progress.add_task(
+                    description=f"Download [green]{escape(path.name)}[/]",
+                    idx=idx,
+                    image=image,
+                    url=link,
+                )
+            else:
+                task_id = progress.add_task(rf"\[[green1]{scan_id}[/]] {escape(path.name)}")
+
+        path.parent.mkdir(exist_ok=True, parents=True)
+        try:
+            downloaded_data = await sandbox.get_file(uri)
+        except SandboxFileNotFoundException:
+            console.warning(f"File {path.name} not found in storage: {uri=} {scan_id=}")
+            if progress:
+                progress.stop_task(task_id)
+                progress.update(task_id=task_id, visible=False)
+            return
+
+        if decompress:
+            try:
+                dctx = zstandard.ZstdDecompressor()
+                with open(path, "wb") as output:
+                    input_fd = BytesIO(downloaded_data)
+                    dctx.copy_stream(input_fd, output)
+            except zstandard.ZstdError as e:
+                console.warning(f"Can't decompress [yellow]{path.name}[/]. {e}")
+        else:
+            async with aiofiles.open(path, "wb") as fd:
+                await fd.write(downloaded_data)
+
+        if progress:
+            progress.stop_task(task_id)
+            progress.update(task_id=task_id, visible=False)
+
+
+async def download(
+    report: SandboxBaseTaskResponse.LongReport,
+    sandbox: Sandbox,
+    out_dir: Path,
+    all: bool = False,
+    artifacts: bool = False,
+    crashdumps: bool = False,
+    debug: bool = False,
+    decompress: bool = False,
+    files: bool = False,
+    logs: bool = False,
+    procdumps: bool = False,
+    video: bool = False,
+    progress: Progress | None = None,
+    idx: str | None = None,
+    image: str | None = None,
+    link: str | None = None,
+) -> None:
+    tasks: list[Coroutine[Any, Any, None]] = []
+
+    def add_task(path: Path, file_uri: str, decompress: bool = False) -> None:
+        tasks.append(
+            _save_artifact(
+                scan_id=report.scan_id,
+                sandbox=sandbox,
+                path=path,
+                file_uri=file_uri,
+                progress=progress,
+                idx=idx,
+                image=image,
+                link=link,
+                decompress=decompress,
+            )
+        )
+
+    for artifact in report.artifacts:
+        sandbox_result = artifact.find_sandbox_result()
+        if sandbox_result is None:
+            continue
+
+        if sandbox_result.details is None:
+            continue
+
+        if sandbox_result.details.sandbox is None:
+            continue
+
+        for log in sandbox_result.details.sandbox.logs:
+            # download default logs by default
+            if (all or logs) and log.type in {
+                LogType.EVENT_CORRELATED,
+                LogType.EVENT_NORMALIZED,
+                LogType.EVENT_RAW,
+                LogType.NETWORK,
+            }:
+                add_task(out_dir / log.file_name, log.file_uri)
+
+            if (all or video) and log.type == LogType.SCREENSHOT:
+                add_task(out_dir / log.file_name, log.file_uri)
+
+            if (all or crashdumps) and log.file_name in {"crashdump.bin", "crashdump.metadata"}:
+                add_task(out_dir / "crashdumps" / log.file_name, log.file_uri)
+
+            if (all or debug) and log.type in {LogType.DEBUG, LogType.GRAPH}:
+                add_task(out_dir / "debug" / log.file_name, log.file_uri)
+
+        if artifacts or files or procdumps or all:
+            if not sandbox_result.details:
+                continue
+
+            if not sandbox_result.details.sandbox:
+                continue
+
+            if not sandbox_result.details.sandbox.artifacts:
+                continue
+
+            for artifact in sandbox_result.details.sandbox.artifacts:
+                if not artifact.file_info:
+                    continue
+
+                if artifact.type == ArtifactType.FILE and (files or artifacts or all):
+                    add_task(
+                        out_dir / "artifacts" / artifact.file_info.file_path.removeprefix("/"),
+                        artifact.file_info.file_uri,
+                    )
+                if artifact.type == ArtifactType.PROCESS_DUMP and (procdumps or artifacts or all):
+                    add_task(
+                        out_dir
+                        / "process_dump"
+                        / artifact.file_info.details.process_dump.process_name.removeprefix("/"),  # type: ignore
+                        artifact.file_info.file_uri,
+                        decompress=decompress,
+                    )
+    if not tasks:
+        console.info(f"Nothing to download from {report.scan_id}")
+
+    await asyncio.gather(*tasks)

sandbox_cli/utils/extractors.py

@@ -0,0 +1,66 @@
+import orjson
+from ptsandbox.models import ArtifactType, EngineSubsystem, SandboxBaseTaskResponse
+
+from sandbox_cli.models.detections import Detections
+
+
+def extract_verdict_from_trace(trace: bytes) -> set[str]:
+    non_generics: set[str] = set()
+    generics: set[str] = set()
+
+    d = Detections(trace)
+    for detect in d.malware:
+        if detect.name.split(".")[-2] == "Generic":
+            generics.add(detect.name)
+        else:
+            non_generics.add(detect.name)
+
+    return non_generics if len(non_generics) != 0 else generics
+
+
+def extract_network_from_trace(trace: bytes) -> set[str]:
+    detects: set[str] = set()
+
+    for line in trace.decode().splitlines(keepends=False):
+        event = orjson.loads(line)
+        if event.get("event.name") == "Auxiliary.ObtainNetworkAlert":
+            detects.add(event.get("s_msg"))
+
+    return detects
+
+
+def extract_static(report: SandboxBaseTaskResponse.LongReport) -> list[str]:
+    sandbox_result = report.artifacts[0].find_sandbox_result()
+
+    ret: set[str] = set()
+
+    if sandbox_result:
+        for artifact in report.artifacts:
+            if artifact.type == ArtifactType.PROCESS_DUMP:
+                continue
+
+            if artifact.engine_results is None:
+                continue
+
+            for artifact_result in artifact.engine_results:
+                if artifact_result.engine_subsystem not in {EngineSubsystem.STATIC, EngineSubsystem.AV}:
+                    continue
+
+                ret |= {f"{artifact_result.engine_code_name}: {x.detect}" for x in artifact_result.detections}
+
+    return sorted(ret)
+
+
+def extract_memory(report: SandboxBaseTaskResponse.LongReport) -> set[str]:
+    ret: set[str] = set()
+
+    sandbox_result = report.artifacts[0].find_sandbox_result()
+
+    if sandbox_result:
+        for artifact in sandbox_result.details.sandbox.artifacts:  # type: ignore
+            if artifact.type == ArtifactType.PROCESS_DUMP:
+                artifact_result = artifact.find_static_result()
+                if artifact_result:
+                    ret |= {x.detect for x in artifact_result.detections}
+
+    return ret

sandbox_cli/utils/merge_dll_hooks.py

@@ -0,0 +1,84 @@
+from pathlib import Path, PurePath
+
+slow_dll_methods = {
+    "RtlAllocateHeap",
+    "GetModuleFileNameA",
+    "GetModuleFileNameW",
+}
+
+
+def __validate_hook(hook: str) -> None:
+    args = hook.split(",")
+
+    if args[0] in ("System32", "SysWOW64"):
+        args = args[1:]
+
+    if len(args) < 1 or not args[0]:
+        raise RuntimeError("Empty function name")
+
+    func_name = args[0]
+    if func_name in slow_dll_methods:
+        raise RuntimeError(f'Function "{func_name}" considered slow')
+
+    if len(args) < 2 or not args[1]:
+        raise RuntimeError("Empty strategy")
+
+    strategy: str
+    function_args: list[str]
+
+    if args[1] == "no-retval":
+        strategy = args[2]
+        function_args = args[3:]
+    else:
+        strategy = args[1]
+        function_args = args[2:]
+
+    if strategy not in ("log", "log+stack"):
+        raise RuntimeError("Bad strategy")
+
+    for arg in function_args:
+        if not arg:
+            raise RuntimeError(f"Extra comma in function argument list of {func_name}")
+
+
+def __merge_dll_hooks(
+    dll_hooks_file: Path,
+    name: PurePath,
+    dll_hooks_cache: set[str],
+) -> bytes:
+    file_hooks = b""
+    with dll_hooks_file.open() as file:
+        for hook in file:
+            hook = hook.strip()
+            if hook == "" or hook.startswith("#"):
+                continue
+
+            __validate_hook(hook)
+            prefix = ""
+            if hook.startswith("SysWOW64"):
+                prefix = "SysWOW64\\"
+            elif hook.startswith("System32"):
+                prefix = "System32\\"
+
+            hook = f"{prefix}{name},{hook}"
+
+            if hook in dll_hooks_cache:
+                continue
+
+            file_hooks += f"{hook}\n".encode()
+            dll_hooks_cache.add(hook)
+    return file_hooks
+
+
+def merge_dll_hooks(dll_hooks_dir: Path) -> bytes:
+    dll_hooks_cache: set[str] = set()
+
+    merged_hooks = b""
+    for hooks_file in dll_hooks_dir.rglob("*.txt"):
+        name = PurePath(hooks_file.relative_to(dll_hooks_dir)).with_suffix(".dll")
+        merged_hooks += __merge_dll_hooks(hooks_file, name, dll_hooks_cache)
+
+    if not merged_hooks:
+        raise RuntimeError("No dll hooks has been found")
+
+    return merged_hooks