[pull] master from kevoreilly:master

tests/test_dist_db.py

@@ -163,3 +163,49 @@ def test_cli_admin_commands():
             modify_node_cli("worker", enabled=False)
             assert mock_node.enabled is False
             mock_print.assert_any_call("Successfully modified node 'worker'.")
+
+
+def test_node_submit_task_error_logging(caplog):
+    import logging
+    from unittest.mock import MagicMock, patch
+    from utils.dist import node_submit_task
+
+    # Setup mock node
+    mock_node = MagicMock()
+    mock_node.id = 42
+    mock_node.name = "worker1"
+    mock_node.url = "http://worker1.local/"
+    mock_node.apikey = "test_key"
+
+    # Setup mock task
+    mock_task = MagicMock()
+    mock_task.id = 16410102
+    mock_task.main_task_id = 3330442
+    mock_task.category = "url"
+    mock_task.path = "http://malicious.url"
+    mock_task.options = {}
+    mock_task.tags = ""
+
+    # Setup mock DB session
+    mock_db = MagicMock()
+    mock_db.scalar.return_value = mock_node
+    mock_db.get.return_value = mock_task
+
+    # Mock response from requests.post
+    mock_response = MagicMock()
+    mock_response.status_code = 503
+    mock_response.text = "Service Unavailable"
+    # To simulate bool(r) is False (requests' Response behavior)
+    mock_response.__bool__.return_value = False
+
+    with patch("utils.dist.requests.post", return_value=mock_response) as mock_post:
+        with caplog.at_level(logging.INFO):
+            result = node_submit_task(task_id=16410102, node_id=42, main_task_id=3330442, db=mock_db)
+
+            assert result is False
+            mock_post.assert_called_once()
+            assert any(
+                "Node: 42 - Task submit to worker failed (CAPE ID: 3330442, Task ID: 16410102): 503 - Service Unavailable"
+                in record.message
+                for record in caplog.records
+            )

utils/dist.py

@@ -607,12 +607,12 @@ def node_submit_task(task_id, node_id, main_task_id, db=None):
             return
 
         # encoding problem
-        if r.status_code == 500 and task.category == "file":
+        if r is not None and r.status_code == 500 and task.category == "file":
             with open(task.path, "rb") as f:
                 r = requests.post(url, data=data, files={"file": ("file", f.read())}, verify=False)
 
         # Zip files preprocessed, so only one id
-        if r and r.status_code == 200:
+        if r is not None and r.status_code == 200:
             if (
                 "task_ids" in r.json().get("data", {})
                 and len(r.json().get("data", {})["task_ids"]) > 0
@@ -637,19 +637,31 @@ def node_submit_task(task_id, node_id, main_task_id, db=None):
             if task.task_id:
                 # log.debug("Submitted task to worker: %s - %d - %d", node.name, task.task_id, task.main_task_id)
                 log.info("[SUBMIT] %-15s <== CAPE ID: %-6d (Worker ID: %d)", node.name, task.main_task_id, task.task_id)
-        elif r and r.status_code == 500:
+        elif r is not None and r.status_code == 500:
             log.info("Saving error to /tmp/dist_error.html")
             _ = path_write_file("/tmp/dist_error.html", r.content)
             log.info((r.status_code, r.text[:200]))
 
-        elif r and r.status_code == 429:
+        elif r is not None and r.status_code == 429:
             log.info((r.status_code, "see api auth for more details"))
 
         else:
-            if r:
-                log.info("Node: %d - Task submit to worker failed: %d - %s", node.id, r.status_code, r.text)
+            if r is not None:
+                log.info(
+                    "Node: %d - Task submit to worker failed (CAPE ID: %d, Task ID: %d): %d - %s",
+                    node.id,
+                    task.main_task_id,
+                    task.id,
+                    r.status_code,
+                    r.text,
+                )
             else:
-                log.info("Node: %d - Task submit to worker failed: No response", node.id)
+                log.info(
+                    "Node: %d - Task submit to worker failed (CAPE ID: %d, Task ID: %d): No response",
+                    node.id,
+                    task.main_task_id,
+                    task.id,
+                )
 
         if check:
             task.node_id = node.id
@@ -816,9 +828,7 @@ def failed_cleaner(self):
                 try:
                     nodes = db.execute(select(Node).where(Node.enabled.is_(True)))
                     for node in nodes or []:
-                        tasks_data = node_fetch_tasks(
-                            "failed_analysis|failed_processing", node.url, node.apikey, action="delete"
-                        )
+                        tasks_data = node_fetch_tasks("failed_analysis|failed_processing", node.url, node.apikey, action="delete")
                         if tasks_data is None:
                             continue
                         failed_task_ids = [task["id"] for task in tasks_data]
@@ -902,7 +912,12 @@ def fetcher(self):
                     if tasks_data is None:
                         self.status_count[node.name] += 1
                         if self.status_count[node.name] >= dead_count:
-                            log.warning("[-] Node %s is unreachable/dead (attempt %d/%d)", node.name, self.status_count[node.name], dead_count)
+                            log.warning(
+                                "[-] Node %s is unreachable/dead (attempt %d/%d)",
+                                node.name,
+                                self.status_count[node.name],
+                                dead_count,
+                            )
                         continue
                     else:
                         self.status_count[node.name] = 0
@@ -1473,7 +1488,13 @@ def submit_tasks(self, node_name, pend_tasks_num, options_like=False, force_push
                                 else:
                                     main_db.set_status(task.main_task_id, TASK_DISTRIBUTED)
                             else:
-                                log.info("something is wrong with submission of task: %d", task.id)
+                                log.info(
+                                    "something is wrong with submission of task %d (CAPE ID: %d) to node %s (Node ID: %d)",
+                                    task.id,
+                                    task.main_task_id,
+                                    node.name,
+                                    node.id,
+                                )
                                 # Since task object is bound to 'db' session which is active in this thread,
                                 # we can use it.
                                 main_db.set_status(task.main_task_id, TASK_BANNED)
@@ -1998,7 +2019,9 @@ class ColoredFormatter(logging.Formatter):
     # Pre-compiled regex patterns (C-level high-performance parsing)
     RE_WORKER = re.compile(r"\b(cape-worker[-\w]*)\b", re.IGNORECASE)
     RE_SETSTAT = re.compile(r"\b(task)\b\s+(\d+)\s+(status)", re.IGNORECASE)
-    RE_MAIN_TASK = re.compile(r"\b(main(?:_task)?_id|main|analysis(?:[\s_-]?id)?|cape(?:[\s_-]?id)?)\b[:\s=]+([\d,]+)", re.IGNORECASE)
+    RE_MAIN_TASK = re.compile(
+        r"\b(main(?:_task)?_id|main|analysis(?:[\s_-]?id)?|cape(?:[\s_-]?id)?)\b[:\s=]+([\d,]+)", re.IGNORECASE
+    )
     RE_WORKER_TASK = re.compile(r"\b(task(?:s|\(s\))?|task_id|worker(?:[\s_-]?id)?)\b[:\s=]+([\d,]+)", re.IGNORECASE)
     RE_PERF = re.compile(r"\bin\s+([0-9.]+s)\b")
 
@@ -2029,34 +2052,19 @@ def format(self, record):
         RESET = "\x1b[0m"
 
         # 1. Highlight all worker/node names dynamically (e.g. "cape-worker", "cape-worker-2")
-        formatted_message = self.RE_WORKER.sub(
-            rf"{CYAN}\1{RESET}",
-            formatted_message
-        )
+        formatted_message = self.RE_WORKER.sub(rf"{CYAN}\1{RESET}", formatted_message)
 
         # 2. Highlight specialized parent/main task updates first (e.g. "task 3319058 status")
-        formatted_message = self.RE_SETSTAT.sub(
-            rf"\1 {MAGENTA}\2{RESET} \3",
-            formatted_message
-        )
+        formatted_message = self.RE_SETSTAT.sub(rf"\1 {MAGENTA}\2{RESET} \3", formatted_message)
 
         # 3. Highlight parent/main task IDs (e.g. "Main ID: 3319058", "Main: 3319058", "main_task_id=3319058", "analysis ID 3319058", "analysis 3319058")
-        formatted_message = self.RE_MAIN_TASK.sub(
-            rf"\1: {MAGENTA}\2{RESET}",
-            formatted_message
-        )
+        formatted_message = self.RE_MAIN_TASK.sub(rf"\1: {MAGENTA}\2{RESET}", formatted_message)
 
         # 4. Highlight general task IDs (e.g. "Task: 16399266", "Task(s): 614534,614541", "task_id=495972")
-        formatted_message = self.RE_WORKER_TASK.sub(
-            rf"\1: {YELLOW}\2{RESET}",
-            formatted_message
-        )
+        formatted_message = self.RE_WORKER_TASK.sub(rf"\1: {YELLOW}\2{RESET}", formatted_message)
 
         # 5. Highlight performance copy benchmarks (e.g. "in 0.08s", "in 0.15s")
-        formatted_message = self.RE_PERF.sub(
-            rf"in {GREEN}\1{RESET}",
-            formatted_message
-        )
+        formatted_message = self.RE_PERF.sub(rf"in {GREEN}\1{RESET}", formatted_message)
 
         # 6. Highlight core operations
         replacements = {
@@ -2327,6 +2335,7 @@ def main():
         log.info("Submit-only mode: Retriever thread disabled.")
 
     import uvicorn
+
     uvicorn.run(app, host=args.host, port=args.port, log_level="debug" if args.debug else "info")