Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for targetrc file to execute commands when a target shell is spawned #859

Merged
merged 8 commits into from
Oct 7, 2024
38 changes: 37 additions & 1 deletion dissect/target/tools/shell.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@

CMD_PREFIX = "cmd_"
_runtime_aliases = {}
DEFAULT_RUNCOMMANDS_FILE = None

def __init__(self, cyber: bool = False):
cmd.Cmd.__init__(self)
Expand All @@ -121,6 +122,28 @@

return object.__getattribute__(self, attr)

def _load_targetrc(self, path: pathlib.Path) -> None:
"""Load and execute commands from the run commands file."""
try:
with path.open() as fh:
for line in fh:
if (line := line.strip()) and not line.startswith("#"): # Ignore empty lines and comments
self.onecmd(line)
except FileNotFoundError:
# The .targetrc file is optional
pass
except Exception as e:
log.debug("Error processing .targetrc file: %s", e)

Check warning on line 136 in dissect/target/tools/shell.py

View check run for this annotation

Codecov / codecov/patch

dissect/target/tools/shell.py#L135-L136

Added lines #L135 - L136 were not covered by tests

def _get_targetrc_path(self) -> pathlib.Path | None:
"""Get the path to the run commands file. Can return ``None`` if ``DEFAULT_RUNCOMMANDS_FILE`` is not set."""
return pathlib.Path(self.DEFAULT_RUNCOMMANDS_FILE).expanduser() if self.DEFAULT_RUNCOMMANDS_FILE else None

Check warning on line 140 in dissect/target/tools/shell.py

View check run for this annotation

Codecov / codecov/patch

dissect/target/tools/shell.py#L140

Added line #L140 was not covered by tests

def preloop(self) -> None:
super().preloop()
if targetrc_path := self._get_targetrc_path():
self._load_targetrc(targetrc_path)

@staticmethod
def check_compatible(target: Target) -> bool:
return True
Expand Down Expand Up @@ -309,6 +332,8 @@
DEFAULT_HISTFILESIZE = 10_000
DEFAULT_HISTDIR = None
DEFAULT_HISTDIRFMT = ".dissect_history_{uid}_{target}"
DEFAULT_RUNCOMMANDS_FILE = "~/.targetrc"
CONFIG_KEY_RUNCOMMANDS_FILE = "TARGETRCFILE"

def __init__(self, target: Target):
self.target = target
Expand Down Expand Up @@ -338,7 +363,15 @@

super().__init__(self.target.props.get("cyber"))

def _get_targetrc_path(self) -> pathlib.Path:
"""Get the path to the run commands file."""

return pathlib.Path(
getattr(self.target._config, self.CONFIG_KEY_RUNCOMMANDS_FILE, self.DEFAULT_RUNCOMMANDS_FILE)
).expanduser()

def preloop(self) -> None:
super().preloop()
if readline and self.histfile.exists():
try:
readline.read_history_file(self.histfile)
Expand Down Expand Up @@ -507,7 +540,6 @@
self.prompt_base = _target_name(target)

TargetCmd.__init__(self, target)

self._clicache = {}
self.cwd = None
self.chdir("/")
Expand Down Expand Up @@ -1144,6 +1176,10 @@
class RegistryCli(TargetCmd):
"""CLI for browsing the registry."""

# Registry shell is incompatible with default shell, so override the default rc file and config key
DEFAULT_RUNCOMMANDS_FILE = "~/.targetrc.registry"
CONFIG_KEY_RUNCOMMANDS_FILE = "TARGETRCFILE_REGISTRY"

def __init__(self, target: Target, registry: regutil.RegfHive | None = None):
self.prompt_base = _target_name(target)

Expand Down
34 changes: 32 additions & 2 deletions tests/tools/test_shell.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import argparse
import pathlib
import platform
import sys
from io import BytesIO, StringIO
from pathlib import Path
from typing import Callable
from unittest.mock import MagicMock
from typing import Callable, Iterator
from unittest.mock import MagicMock, call, mock_open, patch

import pytest

Expand Down Expand Up @@ -130,6 +131,35 @@ def dummy_scandir(path: TargetPath):
assert suggestions == [f"{subfolder_name}/", subfile_name]


@pytest.fixture
def targetrc_file() -> Iterator[list[str]]:
content = """
ls
# This is a comment line and should be ignored
ll
"""

original_open = pathlib.Path.open

def custom_open(self: Path, *args, **kwargs):
if self.name.endswith(".targetrc"):
return mock_open(read_data=content)()
return original_open(self, *args, **kwargs)

with patch("pathlib.Path.open", custom_open):
yield ["ls", "ll"]


def test_targetcli_targetrc(target_bare: Target, targetrc_file: list[str]) -> None:
with patch.object(TargetCli, "onecmd", autospec=True) as mock_onecmd:
cli = TargetCli(target_bare)

cli.preloop()

expected_calls = [call(cli, cmd) for cmd in targetrc_file]
mock_onecmd.assert_has_calls(expected_calls, any_order=False)


@pytest.mark.skipif(platform.system() == "Windows", reason="Unix-specific test.")
def test_pipe_symbol_parsing(capfd: pytest.CaptureFixture, target_bare: Target) -> None:
cli = TargetCli(target_bare)
Expand Down