From fda1ecd37270350a4f757af9e1ae757f1df3b186 Mon Sep 17 00:00:00 2001
From: Thorsten Sick <thorsten.sick@avast.com>
Date: Mon, 14 Jun 2021 13:13:29 +0200
Subject: [PATCH] Sensors can now defined which files are part of the loot

---
 app/experimentcontrol.py                      | 23 +++++++------------
 app/machinecontrol.py                         |  4 +++-
 plugins/base/sensor.py                        |  5 ++--
 .../linux_filebeat/linux_filebeat_plugin.py   |  4 +++-
 tox.ini                                       |  2 +-
 5 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/app/experimentcontrol.py b/app/experimentcontrol.py
index 380ec99..03cb5c6 100644
--- a/app/experimentcontrol.py
+++ b/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):
diff --git a/app/machinecontrol.py b/app/machinecontrol.py
index 5ff7a09..fe24bb3 100644
--- a/app/machinecontrol.py
+++ b/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
 
     ############
 
diff --git a/plugins/base/sensor.py b/plugins/base/sensor.py
index d6469eb..8529048 100644
--- a/plugins/base/sensor.py
+++ b/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
diff --git a/plugins/default/sensors/linux_filebeat/linux_filebeat_plugin.py b/plugins/default/sensors/linux_filebeat/linux_filebeat_plugin.py
index b7ca7b5..b25ac89 100644
--- a/plugins/default/sensors/linux_filebeat/linux_filebeat_plugin.py
+++ b/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]
diff --git a/tox.ini b/tox.ini
index 18a4122..ef5891f 100644
--- a/tox.ini
+++ b/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