dedup do_local_update

CybercentreCanada · Sep 20, 2021 · 1714162 · 1714162
1 parent 2333b82
commit 1714162
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 53 deletions.
diff --git a/suricata_/suricata_.py b/suricata_/suricata_.py
@@ -102,11 +102,14 @@ 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)}")
+        while True:
+            try:
+                # Load the rules
+                self._update_rules()
+                break
+            except Exception as e:
+                raise Exception(f"Something went wrong while trying to load {self.name} rules: {str(e)}")
+            time.sleep(15)
 
         if not self.rules_list:
             self.log.warning("No valid suricata ruleset found. Suricata will run without rules...")

diff --git a/suricata_/update_server.py b/suricata_/update_server.py
@@ -217,54 +217,6 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.updater_type = "suricata"
 
-    def do_local_update(self) -> None:
-        old_update_time = self.get_local_update_time()
-        run_time = time.time()
-        output_directory = tempfile.mkdtemp()
-
-        self.log.info("Setup service account.")
-        username = self.ensure_service_account()
-        self.log.info("Create temporary API key.")
-        with temporary_api_key(self.datastore, username) as api_key:
-            self.log.info(f"Connecting to Assemblyline API: {UI_SERVER}")
-            al_client = get_client(UI_SERVER, apikey=(username, api_key), verify=False)
-
-            # Check if new signatures have been added
-            self.log.info("Check for new signatures.")
-            if al_client.signature.update_available(
-                    since=epoch_to_iso(old_update_time) or '', sig_type=self.updater_type)['update_available']:
-                self.log.info("An update is available for download from the datastore")
-
-                extracted_zip = False
-                attempt = 0
-
-                # Sometimes a zip file isn't always returned, will affect service's use of signature source. Patience..
-                while not extracted_zip and attempt < 5:
-                    temp_zip_file = os.path.join(output_directory, 'temp.zip')
-                    al_client.signature.download(
-                        output=temp_zip_file, query=f"type:{self.updater_type} AND (status:NOISY OR status:DEPLOYED)")
-
-                    if os.path.exists(temp_zip_file):
-                        try:
-                            with ZipFile(temp_zip_file, 'r') as zip_f:
-                                zip_f.extractall(output_directory)
-                                extracted_zip = True
-                                self.log.info("Zip extracted.")
-                        except Exception:
-                            attempt += 1
-                            self.log.warning(f"[{attempt}/5] Bad zip. Trying again after 30s...")
-                            time.sleep(30)
-
-                        os.remove(temp_zip_file)
-
-                if attempt == 5:
-                    self.log.error("Signatures aren't saved to disk. Check sources..")
-                    shutil.rmtree(output_directory, ignore_errors=True)
-                else:
-                    self.log.info("New ruleset successfully downloaded and ready to use")
-                    self.serve_directory(output_directory)
-                    self.set_local_update_time(run_time)
-
     def do_source_update(self, service: Service) -> None:
         self.log.info(f"Connecting to Assemblyline API: {UI_SERVER}...")
         run_time = time.time()