Source updates will come from a PSU

Dockerfile

@@ -94,6 +94,9 @@ RUN touch /etc/suricata/suricata-rules-update
 RUN chown -R assemblyline /var/lib/suricata/
 RUN chown assemblyline /etc/suricata/suricata-rules-update
 
+# Install packages for update-server
+RUN pip install gunicorn flask && rm -rf ~/.cache/pip
+
 # Switch to assemblyline user
 USER assemblyline
 

service_manifest.yml

@@ -51,6 +51,15 @@ docker_config:
   cpu_cores: 1
   ram_mb: 1536
 
+dependencies:
+  updates:
+    container:
+      allow_internet_access: true
+      command: ["python", "-m", "suricata_.update_server"]
+      image: ${REGISTRY}cccs/assemblyline-service-suricata:$SERVICE_TAG
+      ports: ["5003"]
+    run_as_core: True
+
 update_config:
   generates_signatures: true
   method: run

suricata_/suricata_.py

@@ -2,9 +2,13 @@
 import hashlib
 import json
 import os
+import requests
+import shutil
 import subprocess
 import suricatasc
 import sys
+import tarfile
+import tempfile
 import time
 import yaml
 
@@ -20,7 +24,9 @@
 from assemblyline_v4_service.common.result import BODY_FORMAT, Result, ResultSection
 
 SURICATA_BIN = "/usr/local/bin/suricata"
-FILE_UPDATE_DIRECTORY = os.environ.get('FILE_UPDATE_DIRECTORY', '/mount/updates/')
+UPDATES_HOST = os.environ.get('updates_host')
+UPDATES_PORT = os.environ.get('updates_port')
+UPDATES_KEY = os.environ.get('updates_key')
 
 
 class Suricata(ServiceBase):
@@ -37,8 +43,52 @@ def __init__(self, config=None):
         self.suricata_yaml = "/etc/suricata/suricata.yaml"
         self.suricata_log = "/var/log/suricata/suricata.log"
 
-        # Load rules
-        self.rules_hash = self._get_rules_hash()
+        # Updater-related
+        self.rules_directory = None
+        self.update_time = None
+        self.rules_hash = ''
+
+    def _download_rules(self):
+        url_base = f'http://{UPDATES_HOST}:{UPDATES_PORT}/'
+        headers = {
+            'X_APIKEY': UPDATES_KEY
+        }
+
+        # Check if there are new
+        while True:
+            resp = requests.get(url_base + 'status')
+            resp.raise_for_status()
+            status = resp.json()
+            if self.update_time is not None and self.update_time >= status['local_update_time']:
+                return False
+            if status['download_available']:
+                break
+            self.log.warning('Waiting on update server availability...')
+            time.sleep(10)
+
+        # Download the current update
+        temp_directory = tempfile.mkdtemp()
+        buffer_handle, buffer_name = tempfile.mkstemp()
+        try:
+            with os.fdopen(buffer_handle, 'wb') as buffer:
+                resp = requests.get(url_base + 'tar', headers=headers)
+                resp.raise_for_status()
+                for chunk in resp.iter_content(chunk_size=1024):
+                    buffer.write(chunk)
+
+            tar_handle = tarfile.open(buffer_name)
+            tar_handle.extractall(temp_directory)
+            self.update_time = status['local_update_time']
+            self.rules_directory, temp_directory = temp_directory, self.rules_directory
+            return True
+        finally:
+            os.unlink(buffer_name)
+            if temp_directory is not None:
+                shutil.rmtree(temp_directory, ignore_errors=True)
+
+    def _update_rules(self):
+        if self._download_rules():
+            self.rules_hash = self._get_rules_hash()
 
     # Use an external tool to strip frame headers
     @staticmethod
@@ -52,6 +102,12 @@ def strip_frame_headers(filepath):
         return new_filepath
 
     def start(self):
+        try:
+            # Load the rules
+            self._update_rules()
+        except Exception as e:
+            raise Exception(f"Something went wrong while trying to load Suricata rules: {str(e)}")
+
         if not self.rules_list:
             self.log.warning("No valid suricata ruleset found. Suricata will run without rules...")
 
@@ -79,24 +135,16 @@ def start(self):
 
         self.log.info(f"Suricata started with service version: {self.get_service_version()}")
 
-    def _get_rules_hash(self):
-        if not os.path.exists(FILE_UPDATE_DIRECTORY):
-            self.log.warning("Suricata rules directory not found")
-            return None
-
+    def _cleanup(self) -> None:
+        super()._cleanup()
         try:
-            rules_directory = max([os.path.join(FILE_UPDATE_DIRECTORY, d) for d in os.listdir(FILE_UPDATE_DIRECTORY)
-                                   if os.path.isdir(os.path.join(FILE_UPDATE_DIRECTORY, d))
-                                   and not d.startswith('.tmp')],
-                                  key=os.path.getctime)
-        except ValueError:
-            self.log.warning("Suricata rules directory not found")
-            return None
-
-        self.rules_list = [os.path.relpath(str(f), start=FILE_UPDATE_DIRECTORY)
-                           for f in Path(rules_directory).rglob("*") if os.path.isfile(str(f))]
-
-        all_sha256s = [get_sha256_for_file(os.path.join(FILE_UPDATE_DIRECTORY, f)) for f in self.rules_list]
+            self._update_rules()
+        except Exception as e:
+            raise Exception(f"Something went wrong while trying to load Suricata rules: {str(e)}")
+
+    def _get_rules_hash(self):
+        self.rules_list = [str(f) for f in Path(self.rules_directory).rglob("*") if os.path.isfile(str(f))]
+        all_sha256s = [get_sha256_for_file(f) for f in self.rules_list]
 
         self.log.info(f"Suricata will load the following rule files: {self.rules_list}")
 

suricata_/suricata_updater.py

@@ -59,14 +59,14 @@ def url_download(source: Dict[str, Any], previous_update=None) -> List:
 
     LOGGER.info(f"{name} source is configured to {'ignore SSL errors' if ignore_ssl_errors else 'verify SSL'}.")
     if ca_cert:
-        LOGGER.info(f"A CA certificate has been provided with this source.")
+        LOGGER.info("A CA certificate has been provided with this source.")
         add_cacert(ca_cert)
 
     # Create a requests session
     session = requests.Session()
     session.verify = not ignore_ssl_errors
 
-    #Let https requests go through proxy
+    # Let https requests go through proxy
     if proxy:
         os.environ['https_proxy'] = proxy
 
@@ -158,12 +158,12 @@ def git_clone_repo(source: Dict[str, Any], previous_update=None) -> List:
     if ignore_ssl_errors:
         git_env['GIT_SSL_NO_VERIFY'] = 1
 
-    #Let https requests go through proxy
+    # Let https requests go through proxy
     if proxy:
         os.environ['https_proxy'] = proxy
 
     if ca_cert:
-        LOGGER.info(f"A CA certificate has been provided with this source.")
+        LOGGER.info("A CA certificate has been provided with this source.")
         add_cacert(ca_cert)
         git_env['GIT_SSL_CAINFO'] = certifi.where()
 
@@ -199,7 +199,7 @@ def git_clone_repo(source: Dict[str, Any], previous_update=None) -> List:
     else:
         files = [(f, get_sha256_for_file(f)) for f in glob.glob(os.path.join(clone_dir, '*.rules*'))]
 
-    #Clear proxy setting
+    # Clear proxy setting
     if proxy:
         del os.environ['https_proxy']
 
@@ -223,7 +223,7 @@ def suricata_update() -> None:
 
         # Exit if no update sources given
         if 'sources' not in update_config.keys() or not update_config['sources']:
-            LOGGER.error(f"Update configuration does not contain any source to update from")
+            LOGGER.error("Update configuration does not contain any source to update from")
             exit()
 
         # Initialise al_client
@@ -232,7 +232,7 @@ def suricata_update() -> None:
         api_key = update_config['api_key']
         LOGGER.info(f"Connecting to Assemblyline API: {server}...")
         al_client = get_client(server, apikey=(user, api_key), verify=False)
-        LOGGER.info(f"Connected!")
+        LOGGER.info("Connected!")
 
         # Parse updater configuration
         previous_update = update_config.get('previous_update', None)
@@ -296,9 +296,9 @@ def suricata_update() -> None:
             with open(os.path.join(UPDATE_OUTPUT_PATH, 'response.yaml'), 'w') as yml_fh:
                 yaml.safe_dump(dict(hash=json.dumps(files_sha256)), yml_fh)
 
-            LOGGER.info(f"New ruleset successfully downloaded and ready to use")
+            LOGGER.info("New ruleset successfully downloaded and ready to use")
 
-        LOGGER.info(f"Suricata updater completed successfully")
+        LOGGER.info("Suricata updater completed successfully")
     except Exception:
         LOGGER.exception("Updater ended with an exception!")