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

dissect/target/tools/shell.py

@@ -19,7 +19,7 @@
 import sys
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
-from typing import Any, BinaryIO, Callable, Iterator, TextIO
+from typing import Any, BinaryIO, Callable, Iterator, Optional, TextIO
 
 from dissect.cstruct import hexdump
 from flow.record import RecordOutput
@@ -96,6 +96,7 @@
 
     CMD_PREFIX = "cmd_"
     _runtime_aliases = {}
+    DEFAULT_RUNCOMMANDS_FILE = None
 
     def __init__(self, cyber: bool = False):
         cmd.Cmd.__init__(self)
@@ -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)
+
+    def _get_targetrc_path(self) -> Optional[pathlib.Path]:
+        """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
+
+    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
@@ -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
@@ -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)
@@ -507,7 +540,6 @@
         self.prompt_base = _target_name(target)
 
         TargetCmd.__init__(self, target)
-
         self._clicache = {}
         self.cwd = None
         self.chdir("/")
@@ -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)
 

tests/tools/test_shell.py

@@ -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
 
@@ -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)