[pull] master from kevoreilly:master

lib/cuckoo/common/cleaners_utils.py

@@ -18,6 +18,7 @@
 from lib.cuckoo.core.database import Database, _Database
 from lib.cuckoo.core.data.samples import Sample
 from lib.cuckoo.core.data.task import (
+    TASK_BANNED,
     TASK_FAILED_ANALYSIS,
     TASK_FAILED_PROCESSING,
     TASK_FAILED_REPORTING,
@@ -94,6 +95,9 @@ def _close_resolver_pool():
     from dev_utils.elasticsearchdb import all_docs, delete_analysis_and_related_calls, get_analysis_index
 
 
+DB_BATCH_SIZE = 1000
+
+
 def convert_into_time(time_range: str) -> datetime:
     """
     Converts a string representing a time range (e.g., '12h', '1d', '5m')
@@ -218,11 +222,35 @@ def delete_folder(folder):
     @param folder: path to delete.
     @raise CuckooOperationalError: if fails to delete folder.
     """
+    import os
+    import stat
+
+    def remove_readonly(func, path, excinfo):
+        try:
+            # Attempt to heal permissions of the file/dir to make it writable and deletable
+            os.chmod(path, stat.S_IWRITE | stat.S_IREAD)
+            func(path)
+        except Exception:
+            try:
+                # Log detailed ownership and permission information to make debugging easy for the admin
+                st = os.stat(path)
+                import pwd
+                import grp
+                owner = pwd.getpwuid(st.st_uid).pw_name
+                group = grp.getgrgid(st.st_gid).gr_name
+                mode = stat.filemode(st.st_mode)
+                log.error("Failed to delete %s (Mode: %s, Owner: %s, Group: %s)", path, mode, owner, group)
+            except Exception:
+                log.error("Failed to delete %s and failed to fetch ownership details", path)
+
     if path_exists(folder):
         try:
-            shutil.rmtree(folder)
+            shutil.rmtree(folder, onerror=remove_readonly)
+            if path_exists(folder):
+                raise OSError(f"Directory {folder} was not completely removed due to permission errors on child files.")
         except OSError as e:
-            raise CuckooOperationalError(f"Unable to delete folder: {folder}") from e
+            log.error("Unable to delete folder %s: %s", folder, e)
+            raise CuckooOperationalError(f"Unable to delete folder: {folder}. System error: {e}") from e
 
 
 def connect_to_es():
@@ -256,7 +284,19 @@ def is_reporting_db_connected():
         return False
 
 
-def delete_bulk_tasks_n_folders(ids: list, delete_mongo: bool, delete_db_tasks=False):
+def is_contiguous_range(ids: list) -> bool:
+    """Helper to check if a list of IDs is contiguous (has no gaps)."""
+    if not ids:
+        return False
+    sorted_ids = sorted(ids)
+    return (sorted_ids[-1] - sorted_ids[0] + 1) == len(ids)
+
+
+def delete_bulk_tasks_n_folders(ids: list, delete_mongo: bool, delete_db_tasks=False, db_batch_size=None):
+    if db_batch_size is None:
+        db_batch_size = DB_BATCH_SIZE
+
+    # 1. Delete folders sequentially in small batches of 10 to yield disk I/O
     for i in range(0, len(ids), 10):
         ids_tmp = ids[i : i + 10]
         for id in ids_tmp:
@@ -267,17 +307,34 @@ def delete_bulk_tasks_n_folders(ids: list, delete_mongo: bool, delete_db_tasks=F
             except Exception as e:
                 log.error(e)
 
-        if delete_mongo:
-            if mongo_is_cluster():
-                response = input("You are deleting mongo data in cluster, are you sure you want to continue? y/n")
-                if response.lower() in ("n", "not"):
-                    sys.exit()
-            mongo_delete_data(ids_tmp)
-            if delete_db_tasks:
-                try:
-                    db.delete_tasks(task_ids=ids_tmp)
-                except Exception as e:
-                    log.error("Failed to delete tasks from DB: %s", str(e))
+    # 2. Delete from MongoDB in larger, highly-efficient batches or range
+    if delete_mongo and ids:
+        if mongo_is_cluster():
+            response = input("You are deleting mongo data in cluster, are you sure you want to continue? y/n")
+            if response.lower() in ("n", "not"):
+                sys.exit()
+
+        # Check for fast range deletion
+        if is_contiguous_range(ids) and len(ids) > 100:
+            range_start = min(ids)
+            range_end = max(ids) + 1
+            log.info("Contiguous range detected. Deleting MongoDB calls/analyses in range: %d to %d", range_start, range_end)
+            from dev_utils.mongodb import mongo_delete_data_range
+            mongo_delete_data_range(range_start=range_start, range_end=range_end)
+        else:
+            # Otherwise delete in larger batches
+            for i in range(0, len(ids), db_batch_size):
+                ids_batch = ids[i : i + db_batch_size]
+                mongo_delete_data(ids_batch)
+
+    # 3. Delete from SQL database in larger, highly-efficient batches
+    if delete_db_tasks and ids:
+        for i in range(0, len(ids), db_batch_size):
+            ids_batch = ids[i : i + db_batch_size]
+            try:
+                db.delete_tasks(task_ids=ids_batch)
+            except Exception as e:
+                log.error("Failed to delete tasks from DB: %s", str(e))
 
 
 def fail_job(tid):
@@ -407,6 +464,18 @@ def cuckoo_clean_failed_tasks():
     tasks_list = db.delete_tasks(status=f"{TASK_FAILED_ANALYSIS}|{TASK_FAILED_PROCESSING}|{TASK_FAILED_REPORTING}|{TASK_RECOVERED}")
 
 
+def cuckoo_clean_banned_tasks():
+    """Clean up banned tasks
+    It deletes all stored data from file system and configured databases (SQL
+    and MongoDB) for banned tasks.
+    """
+    create_structure()
+
+    tasks_list = db.list_tasks(status=TASK_BANNED)
+    ids = [task.id for task in tasks_list]
+    delete_bulk_tasks_n_folders(ids, delete_mongo=True, delete_db_tasks=True)
+
+
 def cuckoo_clean_bson_suri_logs():
     """Clean up raw suri log files probably not needed if storing in mongo. Does not remove extracted files"""
     # Init logging.
@@ -501,7 +570,7 @@ def tmp_clean_before(timerange: str):
     older_than = convert_into_time(timerange)
     tmp_folder_path = config.cuckoo.get("tmppath")
     # 3rd party?
-    for folder in ("cuckoo-tmp", "cape-external", "cuckoo-sflock"):
+    for folder in ("cuckoo-tmp", "cape-external", "cuckoo-sflock", "cape-pubsub"):
         for root, directories, files in os.walk(os.path.join(tmp_folder_path, folder), topdown=True):
             for name in files + directories:
                 path = os.path.join(root, name)
@@ -802,6 +871,10 @@ def cleanup_files_collection_by_id(task_id: int):
 
 
 def execute_cleanup(args: dict, init_log=True):
+    global DB_BATCH_SIZE
+    if args.get("db_batch_size"):
+        DB_BATCH_SIZE = args["db_batch_size"]
+
     if init_log:
         init_console_logging()
 
@@ -814,6 +887,9 @@ def execute_cleanup(args: dict, init_log=True):
     if args.get("failed_clean"):
         cuckoo_clean_failed_tasks()
 
+    if args.get("banned_clean"):
+        cuckoo_clean_banned_tasks()
+
     if args.get("failed_url_clean"):
         cuckoo_clean_failed_url_tasks()
 

lib/cuckoo/common/web_utils.py

@@ -975,8 +975,9 @@ def download_file(**kwargs):
     generic_demux = False
     if DYNAMIC_PLATFORM_DETERMINATION:
         check_shellcode = "check_shellcode=0" not in kwargs["options"]
+        path_bytes = kwargs["path"] if isinstance(kwargs["path"], bytes) else kwargs["path"].encode()
         _, _, generic_demux = db.identify_submission_package(
-            kwargs["path"].encode(),
+            path_bytes,
             package,
             check_shellcode=check_shellcode,
         )

utils/cleaners.py

@@ -20,6 +20,7 @@
         "--clean", help="Remove all tasks and samples and their associated data", action="store_true", required=False
     )
     parser.add_argument("--failed-clean", help="Remove all tasks marked as failed", action="store_true", required=False)
+    parser.add_argument("--banned-clean", help="Remove all tasks marked as banned", action="store_true", required=False)
     parser.add_argument(
         "--failed-url-clean",
         help="Remove all tasks that are url tasks but we don't have any HTTP traffic",
@@ -105,6 +106,14 @@
         type=int,
     )
 
+    parser.add_argument(
+        "--db-batch-size",
+        help="Batch size for database deletions (MongoDB and SQL). Default is 1000",
+        required=False,
+        default=1000,
+        type=int,
+    )
+
     args = parser.parse_args()
     init_database()
     execute_cleanup(vars(args))