Add support for runtime aliases to target_shell

dissect/target/tools/shell.py

@@ -95,6 +95,7 @@
     """
 
     CMD_PREFIX = "cmd_"
+    _runtime_aliases = {}
 
     def __init__(self, cyber: bool = False):
         cmd.Cmd.__init__(self)
@@ -164,6 +165,11 @@
         return None
 
     def default(self, line: str) -> bool:
+        com, arg, _ = self.parseline(line)
+        if com in self._runtime_aliases:
+            expanded = " ".join([self._runtime_aliases[com], arg])
+            return self.onecmd(expanded)
+
         if (should_exit := self._handle_command(line)) is not None:
             return should_exit
 
@@ -230,6 +236,43 @@
     def complete_man(self, *args) -> list[str]:
         return cmd.Cmd.complete_help(self, *args)
 
+    def do_unalias(self, line: str) -> bool:
+        """delete runtime alias"""
+        aliases = list(shlex.shlex(line, posix=True))
+        for aliased in aliases:
+            if aliased in self._runtime_aliases:
+                del self._runtime_aliases[aliased]
+            else:
+                print(f"alias {aliased} not found")
+        return False
+
+    def do_alias(self, line: str) -> bool:
+        """create a runtime alias"""
+        args = list(shlex.shlex(line, posix=True))
+
+        if not args:
+            for aliased, command in self._runtime_aliases.items():
+                print(f"alias {aliased}={command}")
+            return False
+
+        while args:
+            alias_name = args.pop(0)
+            try:
+                equals = args.pop(0)
+                # our parser works different, so we have to stop this
+                if equals != "=":
+                    raise RuntimeError("Token not allowed")
+                expanded = args.pop(0) if args else ""  # this is how it works in bash
+                self._runtime_aliases[alias_name] = expanded
+            except IndexError:
+                if alias_name in self._runtime_aliases:
+                    print(f"alias {alias_name}={self._runtime_aliases[alias_name]}")
+                else:
+                    print(f"alias {alias_name} not found")
+                pass
+
+        return False
+
     def do_clear(self, line: str) -> bool:
         """clear the terminal screen"""
         os.system("cls||clear")

tests/tools/test_shell.py

@@ -270,3 +270,51 @@ def test_shell_cmd_alias(monkeypatch: pytest.MonkeyPatch, capsys: pytest.Capture
     ls_la_out, _ = run_target_shell(monkeypatch, capsys, target_path, "ls -la")
     ll_out, _ = run_target_shell(monkeypatch, capsys, target_path, "ll")
     assert ls_la_out == ll_out
+
+
+def test_shell_cmd_alias_runtime(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture) -> None:
+    """test if alias commands call their parent attribute correctly."""
+    target_path = absolute_path("_data/tools/info/image.tar")
+
+    # 'list' and 'ls' should return the same output after runtime aliasing
+    list_out, _ = run_target_shell(monkeypatch, capsys, target_path, "alias list=ls xxl='ls -la'\nlist")
+    sys.stdout.flush()
+    ls_out, _ = run_target_shell(monkeypatch, capsys, target_path, "ls")
+    assert list_out == "ubuntu:/$ " + ls_out
+
+    # list aliases
+    sys.stdout.flush()
+    out, _ = run_target_shell(monkeypatch, capsys, target_path, "alias")
+    assert out == "ubuntu:/$ alias list=ls\nalias xxl=ls -la\nubuntu:/$ \n"
+
+    # list single aliases
+    sys.stdout.flush()
+    out, _ = run_target_shell(monkeypatch, capsys, target_path, "alias list")
+    assert out == "ubuntu:/$ alias list=ls\nubuntu:/$ \n"
+
+    # unalias
+    sys.stdout.flush()
+    run_target_shell(monkeypatch, capsys, target_path, "unalias xxl")
+    out, _ = run_target_shell(monkeypatch, capsys, target_path, "alias")
+    assert out == "ubuntu:/$ alias list=ls\nubuntu:/$ \n"
+
+    # unalias multiple and non-existant
+    sys.stdout.flush()
+    out, _ = run_target_shell(monkeypatch, capsys, target_path, "unalias list abc")
+    assert out == "ubuntu:/$ alias abc not found\nubuntu:/$ \n"
+
+    # alias multiple broken - b will be empty
+    sys.stdout.flush()
+    run_target_shell(monkeypatch, capsys, target_path, "alias a=1 b=")
+    out, _ = run_target_shell(monkeypatch, capsys, target_path, "alias")
+    assert out == "ubuntu:/$ alias a=1\nalias b=\nubuntu:/$ \n"
+
+    # alias set/get mixed
+    sys.stdout.flush()
+    out, _ = run_target_shell(monkeypatch, capsys, target_path, "alias b=1 a")
+    assert out == "ubuntu:/$ alias a=1\nubuntu:/$ \n"
+
+    # alias with other symbols not allowed due to parser difference
+    sys.stdout.flush()
+    out, _ = run_target_shell(monkeypatch, capsys, target_path, "alias b+1")
+    assert out.find("*** Unhandled error: Token not allowed") > -1