From 99468b06c084bac41cbd8a34095688e1eefdf023 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Tue, 28 Jan 2025 21:04:45 +0100
Subject: [PATCH] fix: Properly remove pickle tree when indexing fails

---
 bin/background_indexer.py | 4 ++--
 website/web/__init__.py   | 9 +++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/bin/background_indexer.py b/bin/background_indexer.py
index a50839eb..f3265999 100755
--- 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
index 3334e1dd..1f955f23 100644
--- 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))