From a36bd51b6dbb4bf4fe396e75b9b6a4d20334336d Mon Sep 17 00:00:00 2001
From: Miauwkeru <Miauwkeru@users.noreply.github.com>
Date: Mon, 17 Jul 2023 10:45:41 +0000
Subject: [PATCH] Add python test modules and stuff

---
 acquire/acquire.py | 39 ++++++++++++++++++---------------------
 1 file changed, 18 insertions(+), 21 deletions(-)

diff --git a/acquire/acquire.py b/acquire/acquire.py
index 30435ea7..09857073 100644
--- a/acquire/acquire.py
+++ b/acquire/acquire.py
@@ -1329,6 +1329,11 @@ class Boot(Module):
     ]
 
 
+def private_key_filter(path: fsutil.TargetPath) -> bool:
+    with path.open("rt") as file:
+        return "PRIVATE KEY" not in file.readline()
+
+
 @register_module("--home")
 class Home(Module):
     SPEC = [
@@ -1346,33 +1351,25 @@ class Home(Module):
 
 
 @register_module("--ssh")
+@module_arg("--private-keys", action="store_true", help="Add any private keys", default=False)
 class SSH(Module):
-    @classmethod
-    def _run(cls, target: Target, collector):
-        user_pattern = ".ssh/*"
-
-        # Gather user paths
-        # TODO: Use from_user_home if supported for osx
-        if target._os.os == "osx":
-            iterator = [f"/Users/*/{user_pattern}"]
-        else:
-            iterator = list(from_user_home(target, user_pattern))
+    SPEC = [
+        ("glob", ".ssh/*", from_user_home),
+        ("glob", "/etc/ssh/*"),
+        ("glob", "sysvol/ProgramData/ssh/*"),
+    ]
 
+    @classmethod
+    def run(cls, target: Target, cli_args: argparse.Namespace, collector: Collector):
         # Acquire SSH configuration in sshd directories
-        iterator += ["/etc/ssh/*", "sysvol/ProgramData/ssh/*"]
 
-        globbed_path = (path for pattern in iterator for path in target.fs.glob(pattern))
-        for path in globbed_path:
-            if target.fs.path(path).is_dir():
-                collector.collect_dir(path)
-                continue
+        filter = None if cli_args.private_keys else private_key_filter
 
-            with target.fs.path(path).open("rt") as file:
-                if "PRIVATE KEY" in file.readline():
-                    # Detected a private key, skipping.
-                    continue
+        if filter:
+            log.info("Executing SSH without --private-keys, skipping private keys.")
 
-            collector.collect_file(path, outpath=path)
+        with collector.file_filter(filter):
+            super().run(target, cli_args, collector)
 
 
 @register_module("--var")