Agent: Add process list collection PBA

Instead of a system info collector, it is now a PBA.

monkey/infection_monkey/post_breach/actions/collect_processes_list.py

@@ -2,31 +2,33 @@ import logging
 
 import psutil
 
-from common.common_consts.system_info_collectors_names import PROCESS_LIST_COLLECTOR
-from infection_monkey.system_info.system_info_collector import SystemInfoCollector
+from common.common_consts.post_breach_consts import POST_BREACH_PROCESS_LIST_COLLECTION
+from infection_monkey.post_breach.pba import PBA
 
 logger = logging.getLogger(__name__)
 
 # Linux doesn't have WindowsError
+applicable_exceptions = None
 try:
-    WindowsError
+    applicable_exceptions = (psutil.AccessDenied, WindowsError)
 except NameError:
-    # noinspection PyShadowingBuiltins
-    WindowsError = psutil.AccessDenied
+    applicable_exceptions = psutil.AccessDenied
 
 
-class ProcessListCollector(SystemInfoCollector):
+class ProcessListCollection(PBA):
     def __init__(self):
-        super().__init__(name=PROCESS_LIST_COLLECTOR)
+        super().__init__(POST_BREACH_PROCESS_LIST_COLLECTION)
 
-    def collect(self) -> dict:
+    def run(self):
         """
-        Adds process information from the host to the system information.
+        Collects process information from the host.
         Currently lists process name, ID, parent ID, command line
         and the full image path of each process.
         """
         logger.debug("Reading process list")
+
         processes = {}
+        success_state = False
         for process in psutil.process_iter():
             try:
                 processes[process.pid] = {
@@ -36,10 +38,10 @@ class ProcessListCollector(SystemInfoCollector):
                     "cmdline": " ".join(process.cmdline()),
                     "full_image_path": process.exe(),
                 }
-            except (psutil.AccessDenied, WindowsError):
-                # we may be running as non root and some processes are impossible to acquire in
-                # Windows/Linux.
-                # In this case we'll just add what we know.
+                success_state = True
+            except applicable_exceptions:
+                # We may be running as non root and some processes are impossible to acquire in
+                # Windows/Linux. In this case, we'll just add what we know.
                 processes[process.pid] = {
                     "name": "null",
                     "pid": process.pid,
@@ -49,4 +51,4 @@ class ProcessListCollector(SystemInfoCollector):
                 }
                 continue
 
-        return {"process_list": processes}
+        return self.command, [str(processes), success_state]

monkey/infection_monkey/puppet/mock_puppet.py

@@ -12,6 +12,7 @@ from infection_monkey.i_puppet import (
     PortStatus,
     PostBreachData,
 )
+from infection_monkey.post_breach.actions.collect_processes_list import ProcessListCollection
 
 DOT_1 = "10.0.0.1"
 DOT_2 = "10.0.0.2"
@@ -158,6 +159,9 @@ class MockPuppet(IPuppet):
 
         if name == "AccountDiscovery":
             return PostBreachData("pba command 1", ["pba result 1", True])
+        elif name == "ProcessListCollection":
+            cmd, result = ProcessListCollection().run()
+            return PostBreachData(cmd, result)
         else:
             return PostBreachData("pba command 2", ["pba result 2", False])