fix: Properly remove pickle tree when indexing fails

Lookyloo · Jan 28, 2025 · 99468b0 · 99468b0
1 parent 4067a8e
commit 99468b0
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/bin/background_indexer.py b/bin/background_indexer.py
@@ -10,6 +10,7 @@
 
 from lookyloo import Indexing
 from lookyloo.default import AbstractManager, get_config, get_socket_path
+from lookyloo.helpers import remove_pickle_tree
 
 
 logging.config.dictConfig(get_config('logging'))
@@ -53,8 +54,7 @@ def _check_indexes(self) -> None:
                 self.indexing.index_capture(uuid, path)
             except Exception as e:
                 self.logger.error(f'Error while indexing {uuid}: {e}')
-                if (path / 'tree.pickle.gz').exists():
-                    (path / 'tree.pickle.gz').unlink()
+                remove_pickle_tree(path)
             if __counter_shutdown % 10 and self.shutdown_requested():
                 self.logger.warning('Shutdown requested, breaking.')
                 break

diff --git a/website/web/__init__.py b/website/web/__init__.py
@@ -47,7 +47,8 @@
 from lookyloo.helpers import (UserAgents, load_cookies,
                               load_user_config,
                               get_taxonomies,
-                              mimetype_to_generic
+                              mimetype_to_generic,
+                              remove_pickle_tree
                               )
 
 from zoneinfo import available_timezones
@@ -1288,7 +1289,11 @@ def send_mail(tree_uuid: str) -> WerkzeugResponse:
 def trigger_indexing(tree_uuid: str) -> WerkzeugResponse:
     cache = lookyloo.capture_cache(tree_uuid)
     if cache and hasattr(cache, 'capture_dir'):
-        get_indexing(flask_login.current_user).index_capture(tree_uuid, cache.capture_dir)
+        try:
+            get_indexing(flask_login.current_user).index_capture(tree_uuid, cache.capture_dir)
+        except Exception as e:
+            flash(f"Unable to index {tree_uuid}: {e}", 'error')
+            remove_pickle_tree(cache.capture_dir)
     return redirect(url_for('tree', tree_uuid=tree_uuid))