From 665caa324752560d44da9a5aa3011c00014e1c76 Mon Sep 17 00:00:00 2001
From: Steve Garon <steve.garon@cyber.gc.ca>
Date: Fri, 22 May 2020 09:47:49 -0400
Subject: [PATCH] Allow updater to fetch new rules from the database even if
 the sources have no new files

---
 suricata_/suricata_updater.py | 60 ++++++++++++++++++++---------------
 1 file changed, 34 insertions(+), 26 deletions(-)

diff --git a/suricata_/suricata_updater.py b/suricata_/suricata_updater.py
index 12d9f09..67303da 100644
--- a/suricata_/suricata_updater.py
+++ b/suricata_/suricata_updater.py
@@ -175,13 +175,20 @@ 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")
             exit()
 
+        # Initialise al_client
+        server = update_config['ui_server']
+        user = update_config['api_user']
+        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!")
+
         # Parse updater configuration
         previous_update = update_config.get('previous_update', None)
-        previous_hash = update_config.get('previous_hash', None) or {}
-        if previous_hash:
-            previous_hash = json.loads(previous_hash)
+        previous_hash = json.loads(update_config.get('previous_hash', None) or "{}")
         sources = {source['name']: source for source in update_config['sources']}
         files_sha256 = {}
         source_default_classification = {}
@@ -189,7 +196,8 @@ def suricata_update() -> None:
         # Go through each source and download file
         for source_name, source in sources.items():
             uri: str = source['uri']
-            source_default_classification[source_name] = source.get('default_classification', classification.UNRESTRICTED)
+            source_default_classification[source_name] = source.get('default_classification',
+                                                                    classification.UNRESTRICTED)
 
             if uri.endswith('.git'):
                 files = git_clone_repo(source, previous_update=previous_update)
@@ -204,35 +212,31 @@ def suricata_update() -> None:
                     if previous_hash.get(source_name, {}).get(file, None) != sha256:
                         files_sha256[source_name][file] = sha256
 
-        if not files_sha256:
-            LOGGER.info('No Suricata rule file(s) downloaded')
-            shutil.rmtree(UPDATE_OUTPUT_PATH, ignore_errors=True)
-            exit()
-
-        LOGGER.info("Suricata rule(s) file(s) successfully downloaded")
+        if files_sha256:
+            LOGGER.info("Found new Suricata rule files to process!")
 
-        server = update_config['ui_server']
-        user = update_config['api_user']
-        api_key = update_config['api_key']
-        al_client = get_client(server, apikey=(user, api_key), verify=False)
+            suricata_importer = SuricataImporter(al_client, logger=LOGGER)
 
-        suricata_importer = SuricataImporter(al_client, logger=LOGGER)
+            for source, source_val in files_sha256.items():
+                total_imported = 0
+                default_classification = source_default_classification[source]
+                for file in source_val.keys():
+                    total_imported += suricata_importer.import_file(file, source,
+                                                                    default_classification=default_classification)
+                LOGGER.info(f"{total_imported} signatures were imported for source {source}")
 
-        for source, source_val in files_sha256.items():
-            total_imported = 0
-            default_classification = source_default_classification[source]
-            for file in source_val.keys():
-                total_imported += suricata_importer.import_file(file, source, default_classification=default_classification)
-            LOGGER.info(f"{total_imported} signatures were imported for source {source}")
+        else:
+            LOGGER.info('No new Suricata rule files to process')
 
         if al_client.signature.update_available(since=previous_update or '', sig_type='suricata')['update_available']:
-            LOGGER.info("AN UPDATE IS AVAILABLE TO DOWNLOAD")
+            LOGGER.info("An update is available for download from the datastore")
 
             if not os.path.exists(UPDATE_OUTPUT_PATH):
                 os.makedirs(UPDATE_OUTPUT_PATH)
 
             temp_zip_file = os.path.join(UPDATE_OUTPUT_PATH, 'temp.zip')
-            al_client.signature.download(output=temp_zip_file, query="type:suricata AND (status:NOISY OR status:DEPLOYED)")
+            al_client.signature.download(output=temp_zip_file,
+                                         query="type:suricata AND (status:NOISY OR status:DEPLOYED)")
 
             if os.path.exists(temp_zip_file):
                 with ZipFile(temp_zip_file, 'r') as zip_f:
@@ -240,9 +244,13 @@ def suricata_update() -> None:
 
                 os.remove(temp_zip_file)
 
-        # Create the response yaml
-        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)
+            # Create the response yaml
+            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(f"Suricata updater completed successfully")
     except Exception:
         LOGGER.exception("Updater ended with an exception!")