Support loading of runcommand files in registry shell

dissect/target/tools/shell.py

@@ -95,6 +95,7 @@ class ExtendedCmd(cmd.Cmd):
  """
 
  CMD_PREFIX = "cmd_"
+ DEFAULT_RUNCOMMANDSFILE = "~/.targetrc"
 
  def __init__(self, cyber: bool = False):
  cmd.Cmd.__init__(self)
@@ -120,6 +121,28 @@ def print_help(command: str, func: Callable) -> None:
 
  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) -> pathlib.Path:
+ """Get the path to the run commands file."""
+
+ return pathlib.Path(self.DEFAULT_RUNCOMMANDSFILE).expanduser()
+
+ def preloop(self) -> None:
+ super().preloop()
+ self._load_targetrc(self._get_targetrc_path())
+
  @staticmethod
  def check_compatible(target: Target) -> bool:
  return True
@@ -266,6 +289,7 @@ class TargetCmd(ExtendedCmd):
  DEFAULT_HISTFILESIZE = 10_000
  DEFAULT_HISTDIR = None
  DEFAULT_HISTDIRFMT = ".dissect_history_{uid}_{target}"
+ CONFIG_KEY_RUNCOMMANDSFILE = "TARGETRCFILE"
 
  def __init__(self, target: Target):
  self.target = target
@@ -295,7 +319,15 @@ def __init__(self, target: Target):
 
  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_RUNCOMMANDSFILE, self.DEFAULT_RUNCOMMANDSFILE)
+ ).expanduser()
+
  def preloop(self) -> None:
+ super().preloop()
  if readline and self.histfile.exists():
  try:
  readline.read_history_file(self.histfile)
@@ -460,22 +492,14 @@ def do_python(self, line: str) -> bool:
 class TargetCli(TargetCmd):
  """CLI for interacting with a target and browsing the filesystem."""
 
- DEFAULT_RUNCOMMANDSFILE = "~/.targetrc"
-
  def __init__(self, target: Target):
  self.prompt_base = _target_name(target)
 
  TargetCmd.__init__(self, target)
-
  self._clicache = {}
  self.cwd = None
  self.chdir("/")
 
- runcommands_file = pathlib.Path(
- getattr(target._config, "TARGETRCFILE", self.DEFAULT_RUNCOMMANDSFILE)
- ).expanduser()
- self._load_targetrc(runcommands_file)
-
  @property
  def prompt(self) -> str:
  return self.prompt_ps1.format(base=self.prompt_base, cwd=self.cwd)
@@ -494,19 +518,6 @@ def completedefault(self, text: str, line: str, begidx: int, endidx: int) -> lis
  suggestions.append(suggestion)
  return suggestions
 
- 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 resolve_path(self, path: str) -> fsutil.TargetPath:
  if not path:
  return self.cwd
@@ -1121,6 +1132,10 @@ def resolve_glob_path(self, path: fsutil.TargetPath) -> Iterator[fsutil.TargetPa
 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_RUNCOMMANDSFILE = "~/.targetrc.registry"
+ CONFIG_KEY_RUNCOMMANDSFILE = "TARGETRCFILE_REGISTRY"
+
  def __init__(self, target: Target, registry: regutil.RegfHive | None = None):
  self.prompt_base = _target_name(target)
 

tests/tools/test_shell.py

@@ -153,6 +153,9 @@ def custom_open(self: Path, *args, **kwargs):
 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)