chg: [filerwatcher] compression, ext from MH + remove buffer limits

server/workers/workers_2/meta_types_modules/filewatcher/filewatcher.py

@@ -6,15 +6,22 @@ import time
 import os
 import datetime
 import base64
+import shutil
+import gzip
 
 class TypeHandler(MetaTypesDefault):
 
     def __init__(self, uuid, json_file):
         super().__init__(uuid, json_file)
+        if "compress" in json_file:
+            self.compress = json_file['compress']
+        if "extension" in json_file:
+            self.extension = json_file['extension']
         self.set_rotate_file_mode(False)
         self.saved_dir = ''
 
     def process_data(self, data):
+        # Unpack the thing
         self.reconstruct_data(data)
 
     # pushing the filepath instead of the file content to the analyzer
@@ -25,23 +32,44 @@ class TypeHandler(MetaTypesDefault):
 
         # Create folder
         save_dir = os.path.join(self.get_save_dir(save_by_uuid=True), 'files')
-        #debug_dir = os.path.join(self.get_save_dir(), 'debug')
         if not os.path.isdir(save_dir):
             os.makedirs(save_dir)
-        #if not os.path.isdir(debug_dir):
-        #    os.makedirs(debug_dir)
-        # write binary file to disk
+        # write file to disk
         decodeddata = base64.b64decode(data)
+
         m.update(decodeddata)
         path  = os.path.join(save_dir, m.hexdigest())
-        #pathd  = os.path.join(debug_dir, m.hexdigest())
+        path = '{}.{}'.format(path, self.extension)
         with open(path, 'wb') as p:
             p.write(decodeddata)
-
-        #with open(pathd, 'wb') as p:
-        #    p.write(data)
-            # Send data to Analyszer
+        if self.compress:
+            compressed_filename = '{}.gz'.format(path)
+            with open(path, 'rb') as f_in:
+                with gzip.open(compressed_filename, 'wb') as f_out:
+                    shutil.copyfileobj(f_in, f_out)
+            os.remove(path)
+            self.send_to_analyzers(compressed_filename)
+        # Send data to Analyzer
         self.send_to_analyzers(path)
 
-    def test(self):
-        print('Class: filewatcher')
+    def reconstruct_data(self, data):
+        # save data in buffer
+        self.add_to_buffer(data)
+        data = self.get_buffer()
+
+        # end of element found in data
+        if self.get_file_separator() in data:
+            # empty buffer
+            self.reset_buffer()
+            all_line = data.split(self.get_file_separator())
+            for reconstructed_data in all_line[:-1]:
+                if reconstructed_data != b'':
+                    self.handle_reconstructed_data(reconstructed_data)
+
+            # save incomplete element in buffer
+            if all_line[-1] != b'':
+                self.add_to_buffer(all_line[-1])
+
+
+def test(self):
+        print('Class: filewatcher')