Merge pull request #4 from avast/sensors_define_loot

Sensors can now defined which files are part of the loot

app/experimentcontrol.py

@@ -2,7 +2,6 @@
 
 """ A class to control a whole experiment. From setting up the machines to running the attacks """
 
-import glob
 import os
 import subprocess
 import time
@@ -167,9 +166,10 @@ class Experiment():
 
         # Stop sensor plugins
         # Collect data
+        zip_this = []
         for a_target in self.targets:
             a_target.stop_sensors()
-            a_target.collect_sensors(self.lootdir)
+            zip_this += a_target.collect_sensors(self.lootdir)
 
         # Uninstall vulnerabilities
         for a_target in self.targets:
@@ -183,7 +183,7 @@ class Experiment():
         self.__stop_attacker()
 
         self.attack_logger.write_json(os.path.join(self.lootdir, "attack.json"))
-        self.zip_loot()
+        self.zip_loot(zip_this)
 
     def attack(self, target, attack):
         """ Pick an attack and run it
@@ -207,25 +207,18 @@ class Experiment():
             # plugin.__set_logger__(self.attack_logger)
             plugin.__execute__([target])
 
-    def zip_loot(self):
+    def zip_loot(self, zip_this):
         """ Zip the loot together """
 
         filename = os.path.join(self.lootdir, self.starttime + ".zip")
-        globs = ["/**/*.json",
-                 "/**/*.proto",
-                 "/*/**/*.zip",
-
-                 ]
 
         self.attack_logger.vprint(f"Creating zip file {filename}", 1)
 
         with zipfile.ZipFile(filename, "w") as zfh:
-            for a_glob in globs:
-                a_glob = self.lootdir + a_glob
-                for a_file in glob.iglob(a_glob, recursive=True):
-                    if a_file != filename:
-                        self.attack_logger.vprint(a_file, 2)
-                        zfh.write(a_file)
+            for a_file in zip_this:
+                if a_file != filename:
+                    self.attack_logger.vprint(a_file, 2)
+                    zfh.write(a_file)
 
     @staticmethod
     def __get_results_files(root):

app/machinecontrol.py

@@ -251,12 +251,14 @@ class Machine():
 
         machine_specific_path = os.path.join(lootdir, self.config.vmname())
         os.mkdir(machine_specific_path)
+        loot_files = []
 
         for plugin in self.get_sensors():
             self.attack_logger.vprint(f"{CommandlineColors.OKBLUE}Collecting sensor: {plugin.get_name()}{CommandlineColors.ENDC}", 2)
             plugin.set_machine_plugin(self.vm_manager)
-            plugin.__call_collect__(machine_specific_path)
+            loot_files += plugin.__call_collect__(machine_specific_path)
             self.attack_logger.vprint(f"{CommandlineColors.OKGREEN}Collected sensor: {plugin.get_name()}{CommandlineColors.ENDC}", 2)
+        return loot_files
 
     ############
 

plugins/base/sensor.py

@@ -62,11 +62,12 @@ class SensorPlugin(BasePlugin):
 
         path = os.path.join(machine_path, "sensors", self.name)
         os.makedirs(path)
-        self.collect(path)
+        return self.collect(path)
 
-    def collect(self, path):
+    def collect(self, path) -> []:
         """ Collect data from sensor. Copy it from sensor collection dir on target OS to the share
 
         @param path: The path to copy the data into
+        @returns: A list of files to put into the loot zip
         """
         raise NotImplementedError

plugins/default/sensors/linux_filebeat/linux_filebeat_plugin.py

@@ -99,4 +99,6 @@ class LinuxFilebeatPlugin(SensorPlugin):
         """ Collect sensor data """
 
         pg = self.get_playground()
-        self.get_from_machine(f"{pg}/filebeat.json", os.path.join(path, "filebeat.json"))  # nosec
+        dst = os.path.join(path, "filebeat.json")
+        self.get_from_machine(f"{pg}/filebeat.json", dst)  # nosec
+        return [dst]

tox.ini

@@ -4,7 +4,7 @@
 # and then run "tox" from this directory.
 
 [tox]
-envlist = py38
+envlist = py39
 
 [flake8]
 # E501 Line length. Ignored here. But still: please do not abuse the freedom