Make processing info directory configurable

Allow configuration of the path where processing info for failed and bad
samples will be dumped. Having it be statically $HOME/malware_reports
was inflexible.

5 files changed, 26 insertions, 9 deletions

diff --git a/peekaboo.conf.sample b/peekaboo.conf.sample
index 7c9fea1..41830a8 100644
--- a/peekaboo.conf.sample
+++ b/peekaboo.conf.sample
@@ -18,6 +18,7 @@
 #use_debug_module :    no
 # Whether or not to cleanup temporary files under /tmp
 #keep_mail_data   :    no
+#processing_info_dir : /var/lib/peekaboo/malware_reports
 
 
 #[ruleset]
diff --git a/peekaboo/config.py b/peekaboo/config.py
index f921846..9867b12 100644
--- a/peekaboo/config.py
+++ b/peekaboo/config.py
@@ -82,6 +82,7 @@ class PeekabooConfig(object): # pylint: disable=too-many-instance-attributes
         self.job_hash_regex = '/amavis/tmp/([^/]+)/parts/'
         self.use_debug_module = False
         self.keep_mail_data = False
+        self.processing_info_dir = '/var/lib/peekaboo/malware_reports'
         self.db_url = 'sqlite:////var/lib/peekaboo/peekaboo.db'
         self.config_file = '/opt/peekaboo/etc/peekaboo.conf'
         self.ruleset_config = '/opt/peekaboo/etc/ruleset.conf'
@@ -112,6 +113,7 @@ class PeekabooConfig(object): # pylint: disable=too-many-instance-attributes
             'job_hash_regex': ['global', 'job_hash_regex'],
             'use_debug_module': ['global', 'use_debug_module'],
             'keep_mail_data': ['global', 'keep_mail_data'],
+            'processing_info_dir': ['global', 'processing_info_dir'],
             'db_url': ['db', 'url'],
             'ruleset_config': ['ruleset', 'config'],
             'cuckoo_mode': ['cuckoo', 'mode'],
diff --git a/peekaboo/daemon.py b/peekaboo/daemon.py
index eaf007c..326ea03 100644
--- a/peekaboo/daemon.py
+++ b/peekaboo/daemon.py
@@ -328,7 +328,7 @@ def run():
     # database connection and connection map.
     sample_factory = SampleFactory(
         cuckoo, config.sample_base_dir, config.job_hash_regex,
-        config.keep_mail_data)
+        config.keep_mail_data, config.processing_info_dir)
 
     # We only want to accept 2 * worker_count connections.
     try:
diff --git a/peekaboo/sample.py b/peekaboo/sample.py
index fe79b55..34b4a57 100644
--- a/peekaboo/sample.py
+++ b/peekaboo/sample.py
@@ -47,7 +47,7 @@ class SampleFactory(object):
     sample needs and thus serves as a registry of potential API breakage
     perhaps deserving looking into. """
     def __init__(self, cuckoo, base_dir, job_hash_regex,
-                 keep_mail_data):
+                 keep_mail_data, processing_info_dir):
         # object references for interaction
         self.cuckoo = cuckoo
 
@@ -55,12 +55,14 @@ class SampleFactory(object):
         self.base_dir = base_dir
         self.job_hash_regex = job_hash_regex
         self.keep_mail_data = keep_mail_data
+        self.processing_info_dir = processing_info_dir
 
     def make_sample(self, file_path, status_change=None, metainfo=None):
         """ Create a new Sample object based on the factory's configured
         defaults and variable parameters. """
         return Sample(file_path, self.cuckoo, status_change, metainfo,
-                      self.base_dir, self.job_hash_regex, self.keep_mail_data)
+                      self.base_dir, self.job_hash_regex, self.keep_mail_data,
+                      self.processing_info_dir)
 
 
 class Sample(object):
@@ -79,7 +81,7 @@ class Sample(object):
     """
     def __init__(self, file_path, cuckoo=None, status_change=None,
                  metainfo=None, base_dir=None, job_hash_regex=None,
-                 keep_mail_data=False):
+                 keep_mail_data=False, processing_info_dir=None):
         self.__path = file_path
         self.__cuckoo = cuckoo
         self.__wd = None
@@ -107,6 +109,7 @@ class Sample(object):
         self.__job_hash = None
         self.__job_hash_regex = job_hash_regex
         self.__keep_mail_data = keep_mail_data
+        self.__processing_info_dir = processing_info_dir
         self.initialized = False
 
         if metainfo:
@@ -325,8 +328,12 @@ class Sample(object):
         Saves the Cuckoo report as HTML + JSON
         to a directory named after the job hash.
         """
-        dump_dir = os.path.join(os.environ['HOME'], 'malware_reports',
-                                self.job_hash)
+        if not self.__processing_info_dir:
+            logger.debug('Not dumping processing info because no path for the '
+                         'data is unconfigured.')
+            return
+
+        dump_dir = os.path.join(self.__processing_info_dir, self.job_hash)
         if not os.path.isdir(dump_dir):
             os.makedirs(dump_dir, 0o770)
         filename = self.__filename + '-' + self.sha256sum
diff --git a/test.py b/test.py
index 9d9e72c..bf765fe 100644
--- a/test.py
+++ b/test.py
@@ -89,6 +89,9 @@ class TestDefaultConfig(TestConfig):
         self.assertEqual(self.config.use_debug_module, False)
         self.assertEqual(self.config.keep_mail_data, False)
         self.assertEqual(
+            self.config.processing_info_dir,
+            '/var/lib/peekaboo/malware_reports')
+        self.assertEqual(
             self.config.ruleset_config, '/opt/peekaboo/etc/ruleset.conf')
         self.assertEqual(self.config.log_level, logging.INFO)
         self.assertEqual(
@@ -119,6 +122,7 @@ sample_base_dir  :    /tmp/1
 job_hash_regex   :    /var/2
 use_debug_module :    yes
 keep_mail_data   :    yes
+processing_info_dir : /var/3
 
 [ruleset]
 config           :    /rules/1
@@ -157,6 +161,7 @@ duplicate_check_interval: 61
         self.assertEqual(self.config.job_hash_regex, '/var/2')
         self.assertEqual(self.config.use_debug_module, True)
         self.assertEqual(self.config.keep_mail_data, True)
+        self.assertEqual(self.config.processing_info_dir, '/var/3')
         self.assertEqual(self.config.ruleset_config, '/rules/1')
         self.assertEqual(self.config.log_level, logging.DEBUG)
         self.assertEqual(self.config.log_format, 'format%foo1')
@@ -359,7 +364,8 @@ class TestDatabase(unittest.TestCase):
                                              instance_id=0)
         cls.factory = CreatingSampleFactory(
             cuckoo=None, base_dir=cls.conf.sample_base_dir,
-            job_hash_regex=cls.conf.job_hash_regex, keep_mail_data=False)
+            job_hash_regex=cls.conf.job_hash_regex, keep_mail_data=False,
+            processing_info_dir=None)
         cls.sample = cls.factory.create_sample('test.py', 'test')
         result = RuleResult('Unittest',
                             Result.failed,
@@ -503,7 +509,8 @@ class TestSample(unittest.TestCase):
         cls.db_con = PeekabooDatabase('sqlite:///' + cls.test_db)
         cls.factory = CreatingSampleFactory(
             cuckoo=None, base_dir=cls.conf.sample_base_dir,
-            job_hash_regex=cls.conf.job_hash_regex, keep_mail_data=False)
+            job_hash_regex=cls.conf.job_hash_regex, keep_mail_data=False,
+            processing_info_dir=None)
         cls.sample = cls.factory.create_sample('test.py', 'test')
 
     def test_attribute_dict(self):
@@ -531,7 +538,7 @@ class TestSample(unittest.TestCase):
         legacy_factory = CreatingSampleFactory(
             cuckoo=None, base_dir=self.conf.sample_base_dir,
             job_hash_regex=r'/var/lib/amavis/tmp/([^/]+)/parts.*',
-            keep_mail_data=False)
+            keep_mail_data=False, processing_info_dir=None)
         sample = legacy_factory.make_sample(path_with_job_hash, 'file')
         self.assertEqual(job_hash, sample.job_hash,
                          'Job hash regex is not working')