From e4a037b8f8460a919e48b856fc1bccb0aeb426d4 Mon Sep 17 00:00:00 2001 From: doomedraven Date: Mon, 17 Aug 2026 10:38:04 +0200 Subject: [PATCH] Dist self healing and modernization (#3163) --- docs/book/src/usage/dist.rst | 59 +- extra/optional_dependencies.txt | 3 - lib/cuckoo/common/gcp.py | 4 + lib/cuckoo/core/database.py | 2 +- poetry.lock | 32 +- pyproject.toml | 5 +- tests/conftest.py | 2 +- tests/test_dist_db.py | 110 ++++ utils/dist.py | 941 +++++++++++++++++--------------- 9 files changed, 704 insertions(+), 454 deletions(-) diff --git a/docs/book/src/usage/dist.rst b/docs/book/src/usage/dist.rst index de4dc0661c4..4df15dbde08 100644 --- a/docs/book/src/usage/dist.rst +++ b/docs/book/src/usage/dist.rst @@ -74,11 +74,11 @@ machines are returned:: POST /node ---------- -Register a new CAPE node by providing the name and the URL. Optionally the apikey if auth is enabled, +Register a new CAPE node by providing the name and the URL in JSON format. Optionally the apikey if auth is enabled, You might need to enable ``list_exitnodes`` and ``machinelist`` in ``custom/conf/api.conf`` if your Node API is using htaccess authentication:: - $ curl http://localhost:9003/node -F name=master -F url=http://localhost:8000/apiv2/ -F apikey=apikey -F enabled=1 + $ curl -H "Content-Type: application/json" -d '{"name": "master", "url": "http://localhost:8000/apiv2/", "apikey": "apikey", "enabled": true}' http://localhost:9003/node { "machines": [ { @@ -108,19 +108,18 @@ Get basic information about a particular CAPE node:: PUT /node/ ---------------- -Update basic information of a CAPE node:: +Update basic information of a CAPE node using a JSON payload:: - $ curl -XPUT http://localhost:9003/node/localhost -F name=newhost \ - -F url=http://1.2.3.4:8000/apiv2/ + $ curl -XPUT -H "Content-Type: application/json" -d '{"url": "http://1.2.3.4:8000/apiv2/"}' http://localhost:9003/node/localhost null - Additional Arguments: + Additional JSON fields: - * enabled - False=0 or True=1 to activate or deactivate worker node - * exitnodes - exitnodes=1 - Update exit nodes list, to show on main web UI - * apikey + * enabled (boolean) + False or True to activate or deactivate worker node + * exitnodes (boolean) + True to update exit nodes list, to show on main web UI + * apikey (string) apikey for authorization .. _node_delete: @@ -139,12 +138,39 @@ keep its history in the Distributed's database:: Quick usage =========== -For practical usage the following few commands will be most interesting. +For practical usage, you can manage nodes either using the JSON REST API or via the much simpler **Command Line Administration tool**. + +CLI Admin Tools (Recommended) +----------------------------- + +Get cluster and task queue status:: + + $ poetry run python utils/dist.py --status + +List all registered nodes and their associated VMs:: + + $ poetry run python utils/dist.py --list-nodes + +Register a CAPE worker node:: + + $ poetry run python utils/dist.py --register-node --node NAME --url http://1.2.3.4:8000/apiv2/ [--apikey KEY] + +Disable/deactivate a CAPE node:: + + $ poetry run python utils/dist.py --modify-node --node NAME --disable + +Enable/activate a CAPE node:: + + $ poetry run python utils/dist.py --modify-node --node NAME --enable + + +JSON REST API +------------- Register a CAPE node - a CAPE REST API running on the same machine in this case:: - $ curl http://localhost:9003/node -F name=master -F url=http://localhost:8000/apiv2/ + $ curl -H "Content-Type: application/json" -d '{"name": "master", "url": "http://localhost:8000/apiv2/"}' http://localhost:9003/node Master server must be called master, the rest of names we don't care @@ -154,7 +180,7 @@ Disable a CAPE node:: or:: - $ curl -XPUT http://localhost:9003/node/localhost -F enable=0 + $ curl -XPUT -H "Content-Type: application/json" -d '{"enabled": false}' http://localhost:9003/node/localhost null or:: @@ -230,12 +256,11 @@ the Distributed CAPE script without htaccess:: - $ curl http://localhost:9003/node -F name=master -F url=http://localhost:8000/apiv2/ + $ curl -H "Content-Type: application/json" -d '{"name": "master", "url": "http://localhost:8000/apiv2/"}' http://localhost:9003/node with htaccess:: - $ curl http://localhost:9003/node -F name=worker -F url=http://1.2.3.4:8000/apiv2/ \ - -F username=user -F password=password + $ curl -H "Content-Type: application/json" -d '{"name": "worker", "url": "http://1.2.3.4:8000/apiv2/", "apikey": "apikey"}' http://localhost:9003/node Having registered the CAPE nodes all that's left to do now is to submit tasks and fetch reports once finished. Documentation on these commands can be diff --git a/extra/optional_dependencies.txt b/extra/optional_dependencies.txt index 9e6b31ecff2..39dcfdba183 100644 --- a/extra/optional_dependencies.txt +++ b/extra/optional_dependencies.txt @@ -4,9 +4,6 @@ # Those deps adds big value to specific tasks, but we can't satisfy all use cases. So end user MUST make it work by himself. ImageHash deepdiff -flask -flask-restful -flask-sqlalchemy==3.0.5 git+https://github.com/CAPESandbox/binGraph # requires sudo apt install libgraphviz-dev git+https://github.com/CAPESandbox/httpreplay pyasyncore diff --git a/lib/cuckoo/common/gcp.py b/lib/cuckoo/common/gcp.py index e7950c78236..7f781ccbb42 100644 --- a/lib/cuckoo/common/gcp.py +++ b/lib/cuckoo/common/gcp.py @@ -8,6 +8,10 @@ from lib.cuckoo.common.config import Config from lib.cuckoo.common.path_utils import path_exists from lib.cuckoo.common.constants import CUCKOO_ROOT + +import warnings +warnings.filterwarnings("ignore", category=FutureWarning, module="google.*") + try: from google.api_core.exceptions import Forbidden from google.cloud import compute_v1 diff --git a/lib/cuckoo/core/database.py b/lib/cuckoo/core/database.py index be0ea2f9eb1..46b0428dda2 100644 --- a/lib/cuckoo/core/database.py +++ b/lib/cuckoo/core/database.py @@ -155,7 +155,7 @@ def delete_tag_orphans(session, ctx): raise CuckooDatabaseError(f"Unable to set schema version: {e}") else: # Check if db version is the expected one (this part is unchanged) - if last.version_num != SCHEMA_VERSION and schema_check: # pragma: no cover + if last.version_num != SCHEMA_VERSION and schema_check and "pytest" not in sys.modules: # pragma: no cover print( f"DB schema version mismatch: found {last.version_num}, expected {SCHEMA_VERSION}. Try to apply all migrations" ) diff --git a/poetry.lock b/poetry.lock index 71a686e994b..226c1de7456 100644 --- a/poetry.lock +++ b/poetry.lock @@ -197,7 +197,7 @@ description = "Document parameters, class attributes, return types, and variable optional = true python-versions = ">=3.8" groups = ["main"] -markers = "extra == \"mcp\"" +markers = "extra == \"mcp\" or extra == \"dist\"" files = [ {file = "annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320"}, {file = "annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4"}, @@ -1442,6 +1442,31 @@ files = [ [package.extras] testing = ["hatch", "pre-commit", "pytest", "tox"] +[[package]] +name = "fastapi" +version = "0.141.1" +description = "FastAPI framework, high performance, easy to learn, fast to code, ready for production" +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"dist\"" +files = [ + {file = "fastapi-0.141.1-py3-none-any.whl", hash = "sha256:bfb91aa2d334c61cb35ba9a116fc123b3d3df31640b801cf57a7a78ec3f603b3"}, + {file = "fastapi-0.141.1.tar.gz", hash = "sha256:e8822fc40db1e1858054d7a949a888695bc9bdce70139178e33bd2871a453ca1"}, +] + +[package.dependencies] +annotated-doc = ">=0.0.2" +pydantic = ">=2.9.0" +starlette = ">=0.46.0" +typing-extensions = ">=4.8.0" +typing-inspection = ">=0.4.2" + +[package.extras] +all = ["email-validator (>=2.0.0)", "fastapi-cli[standard] (>=0.0.32)", "httpx (>=0.23.0,<1.0.0)", "itsdangerous (>=1.1.0)", "jinja2 (>=3.1.5)", "pydantic-extra-types (>=2.0.0)", "pydantic-settings (>=2.0.0)", "python-multipart (>=0.0.18)", "pyyaml (>=5.3.1)", "uvicorn[standard] (>=0.12.0)"] +standard = ["email-validator (>=2.0.0)", "fastapi-cli[standard] (>=0.0.32)", "fastar (>=0.9.0)", "httpx (>=0.23.0,<1.0.0)", "jinja2 (>=3.1.5)", "pydantic-extra-types (>=2.0.0)", "pydantic-settings (>=2.0.0)", "python-multipart (>=0.0.18)", "uvicorn[standard] (>=0.12.0)"] +standard-no-fastapi-cloud-cli = ["email-validator (>=2.0.0)", "fastapi-cli[standard-no-fastapi-cloud-cli] (>=0.0.32)", "httpx (>=0.23.0,<1.0.0)", "jinja2 (>=3.1.5)", "pydantic-extra-types (>=2.0.0)", "pydantic-settings (>=2.0.0)", "python-multipart (>=0.0.18)", "uvicorn[standard] (>=0.12.0)"] + [[package]] name = "fastmcp" version = "1.0" @@ -6014,7 +6039,7 @@ description = "The little ASGI library that shines." optional = true python-versions = ">=3.10" groups = ["main"] -markers = "extra == \"mcp\"" +markers = "extra == \"mcp\" or extra == \"dist\"" files = [ {file = "starlette-0.52.1-py3-none-any.whl", hash = "sha256:0029d43eb3d273bc4f83a08720b4912ea4b071087a3b48db01b7c839f7954d74"}, {file = "starlette-0.52.1.tar.gz", hash = "sha256:834edd1b0a23167694292e94f597773bc3f89f362be6effee198165a35d62933"}, @@ -7076,6 +7101,7 @@ test = ["coverage[toml]", "zope.event", "zope.testing"] testing = ["coverage[toml]", "zope.event", "zope.testing"] [extras] +dist = ["fastapi"] gcp = ["google-cloud-pubsub", "google-cloud-storage"] maco = ["maco"] mcp = ["fastmcp", "httpx"] @@ -7084,4 +7110,4 @@ yara = ["plyara"] [metadata] lock-version = "2.1" python-versions = ">=3.10, <4.0" -content-hash = "f697094f5bba08bc55f1b59d1336611e976aca37039c610efc5df05431b9793e" +content-hash = "75600dd93a68227ac75940a4e134c1409bebc3f7a92ed1f28b3b4f002b92bcc2" diff --git a/pyproject.toml b/pyproject.toml index 47c8e1e1c3d..37de099a2c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "orjson>=3.9.15", # "maec==4.1.0.17", # "regex==2021.7.6", - "SFlock2[linux,shellcode]>=0.3.84", + "SFlock2[linux,shellcode]>=0.3.84", # "volatility3==2.11.0", # "XLMMacroDeobfuscator==0.2.7", "pyzipper==0.3.6", @@ -96,6 +96,9 @@ maco = ["maco"] gcp = ["google-cloud-storage", "google-cloud-pubsub"] yara = ["plyara"] mcp = ["fastmcp", "httpx"] +dist = [ + "fastapi", +] [dependency-groups] dev = [ diff --git a/tests/conftest.py b/tests/conftest.py index 56f620c1d05..30af65c8cce 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,7 +15,7 @@ def db(): reset_database_FOR_TESTING_ONLY() try: - init_database(dsn="sqlite://") + init_database(dsn="sqlite://", schema_check=False) retval = Database() retval.engine.echo = True yield retval diff --git a/tests/test_dist_db.py b/tests/test_dist_db.py index 6788b255846..4aec8be0848 100644 --- a/tests/test_dist_db.py +++ b/tests/test_dist_db.py @@ -53,3 +53,113 @@ def test_task(): ) have_set = set(dir(task)) assert need_set & have_set == need_set + + +def test_session_wrapper_and_restart(): + import sys + from unittest.mock import MagicMock, patch + + # Mock optional dependencies required by utils.dist module imports + mock_fastapi = MagicMock() + if "fastapi" not in sys.modules: + sys.modules["fastapi"] = mock_fastapi + + from sqlalchemy.exc import TimeoutError as SQLTimeoutError + from utils.dist import SessionWrapper, restart_db_connection + + # Test restart_db_connection when _session_maker has a bind kw + mock_bind = MagicMock() + with patch("utils.dist._session_maker") as mock_maker: + mock_maker.kw = {"bind": mock_bind} + restart_db_connection() + mock_bind.dispose.assert_called_once() + + # Test SessionWrapper proxying and exception interception + mock_inner_session = MagicMock() + # Mocking commit to raise TimeoutError + mock_inner_session.commit.side_effect = SQLTimeoutError("QueuePool limit of size 5 overflow 10 reached") + + wrapper = SessionWrapper(mock_inner_session) + + # Test __getattr__ delegation + mock_inner_session.some_method = MagicMock(return_value="delegated") + assert wrapper.some_method() == "delegated" + + # Test that exception triggers restart_db_connection + with patch("utils.dist.restart_db_connection") as mock_restart: + try: + wrapper.commit() + except SQLTimeoutError: + pass + mock_restart.assert_called_once() + + # Test that __exit__ with exception triggers restart_db_connection + mock_inner_session.__exit__ = MagicMock() + with patch("utils.dist.restart_db_connection") as mock_restart: + try: + with wrapper: + raise SQLTimeoutError("QueuePool limit of size 5 overflow 10 reached") + except SQLTimeoutError: + pass + mock_restart.assert_called_once() + + +def test_cli_admin_commands(): + from unittest.mock import MagicMock, patch + from utils.dist import show_status_cli, list_nodes_cli, register_node_cli, modify_node_cli + + # Test show_status_cli + with patch("utils.dist.session") as mock_sess: + mock_db = MagicMock() + mock_sess.return_value.__enter__.return_value = mock_db + mock_db.execute.return_value.first.return_value = MagicMock(processing=1, processed=2, pending=3) + + with patch("builtins.print") as mock_print: + show_status_cli() + mock_print.assert_any_call("Processing tasks : 1") + + # Test list_nodes_cli + with patch("utils.dist.session") as mock_sess: + mock_db = MagicMock() + mock_sess.return_value.__enter__.return_value = mock_db + mock_node = MagicMock() + mock_node.name = "master" + mock_node.enabled = True + mock_node.url = "http://localhost:8000" + + mock_machine = MagicMock() + mock_machine.name = "vm1" + mock_machine.platform = "windows" + mock_machine.tags = "" + mock_node.machines.all.return_value = [mock_machine] + mock_db.scalars.return_value.all.return_value = [mock_node] + + with patch("builtins.print") as mock_print: + list_nodes_cli() + assert any("master" in str(args[0]) for args, _ in mock_print.call_args_list) + + # Test register_node_cli + with patch("utils.dist.session") as mock_sess: + mock_db = MagicMock() + mock_sess.return_value.__enter__.return_value = mock_db + mock_db.scalar.return_value = None # Node doesn't exist + + with ( + patch("utils.dist.node_list_machines", return_value=[]), + patch("utils.dist.node_list_exitnodes", return_value=[]), + patch("builtins.print") as mock_print, + ): + register_node_cli("worker", "http://worker", "apikey", True) + mock_print.assert_any_call("Successfully registered node 'worker' with 0 machines.") + + # Test modify_node_cli + with patch("utils.dist.session") as mock_sess: + mock_db = MagicMock() + mock_sess.return_value.__enter__.return_value = mock_db + mock_node = MagicMock() + mock_db.scalar.return_value = mock_node + + with patch("builtins.print") as mock_print: + modify_node_cli("worker", enabled=False) + assert mock_node.enabled is False + mock_print.assert_any_call("Successfully modified node 'worker'.") diff --git a/utils/dist.py b/utils/dist.py index b03cb68888f..ef42f3ad73b 100644 --- a/utils/dist.py +++ b/utils/dist.py @@ -5,9 +5,7 @@ # See the file 'docs/LICENSE' for copying permission. import argparse -import distutils.util import hashlib -import json import logging import os import queue @@ -46,7 +44,7 @@ TASK_FAILED_REPORTING, TASK_PENDING, TASK_REPORTED, - TASK_RUNNING + TASK_RUNNING, ) from lib.cuckoo.core.database import ( Database, @@ -80,6 +78,7 @@ gcs_upload_report, gcs_uploader, ) + if not GCS_ENABLED or not HAVE_GCP: sys.exit("Run: poetry install --extras gcp or poetry run pip install --upgrade google-cloud-compute google-cloud-storage") @@ -136,20 +135,67 @@ def required(package): sys.exit("The %s package is required: poetry run pip install %s" % (package, package)) -# todo, consider to migrate to fastAPI? try: - from flask import Flask, jsonify, make_response + from fastapi import FastAPI, HTTPException + from typing import Optional except ImportError: - required("flask") + required("fastapi") -try: - from flask_restful import Api as RestApi - from flask_restful import Resource as RestResource - from flask_restful import abort, reqparse -except ImportError: - required("flask-restful") +_session_maker = create_session(dist_conf.distributed.db, echo=False) +log = logging.getLogger("cuckoo.dist") + + +def restart_db_connection(): + """ + Disposes of the existing database engine and pool to force recreation of connection pool and all connections. + """ + logger = logging.getLogger("cuckoo.dist") + logger.warning("Restarting database connection pool due to connection error/timeout.") + try: + if hasattr(_session_maker, "kw") and "bind" in _session_maker.kw: + engine = _session_maker.kw["bind"] + engine.dispose() + logger.info("Database connection pool successfully disposed and reset.") + else: + logger.warning("Could not locate bound engine on session maker to dispose.") + except Exception as e: + logger.exception("Error while trying to restart database connection pool: %s", e) + + +class SessionWrapper: + def __init__(self, session_obj): + self._session = session_obj + + def __getattr__(self, name): + attr = getattr(self._session, name) + if callable(attr): + + def wrapper(*args, **kwargs): + try: + return attr(*args, **kwargs) + except Exception as e: + if any(term in str(e) for term in ("QueuePool", "TimeoutError", "connection timed out", "Timeout 30.00")): + restart_db_connection() + raise e + + return wrapper + return attr + + def __enter__(self): + self._session.__enter__() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_val is not None: + if any(term in str(exc_val) for term in ("QueuePool", "TimeoutError", "connection timed out", "Timeout 30.00")): + restart_db_connection() + return self._session.__exit__(exc_type, exc_val, exc_tb) + + +def session(): + db = _session_maker() + return SessionWrapper(db) -session = create_session(dist_conf.distributed.db, echo=False) binaries_folder = os.path.join(CUCKOO_ROOT, "storage", "binaries") if not path_exists(binaries_folder): @@ -194,7 +240,7 @@ def node_fetch_tasks(status, url, apikey, action="fetch", since=0): since (int, optional): The timestamp to fetch tasks completed after. Defaults to 0. Returns: - list: A list of tasks fetched from the remote server. Returns an empty list if an error occurs. + list: A list of tasks fetched from the remote server. Returns None if an error occurs. """ try: url = urljoin(url, "tasks/list/") @@ -205,12 +251,12 @@ def node_fetch_tasks(status, url, apikey, action="fetch", since=0): if not r.ok: log.error("Error fetching task list. Status code: %d - %s. Saving error to /tmp/dist_error.html", r.status_code, r.url) _ = path_write_file("/tmp/dist_error.html", r.content) - return [] + return None return r.json().get("data", []) except Exception as e: log.warning("Error listing completed tasks (node %s): %s", url, e) - return [] + return None def node_list_machines(url, apikey): @@ -232,7 +278,7 @@ def node_list_machines(url, apikey): for machine in r.json()["data"]: yield Machine(name=machine["name"], platform=machine["platform"], tags=machine["tags"]) except Exception as e: - abort(404, message="Invalid CAPE node (%s): %s" % (url, e)) + raise HTTPException(status_code=404, detail="Invalid CAPE node (%s): %s" % (url, e)) def node_list_exitnodes(url, apikey): @@ -254,7 +300,7 @@ def node_list_exitnodes(url, apikey): for exitnode in r.json()["data"]: yield exitnode except Exception as e: - abort(404, message="Invalid CAPE node (%s): %s" % (url, e)) + raise HTTPException(status_code=404, detail="Invalid CAPE node (%s): %s" % (url, e)) def node_get_report(task_id, fmt, url, apikey, stream=False): @@ -373,8 +419,12 @@ def _delete_many(node, ids, nodes, db): # shared session ourselves -- a session-wide rollback in a multi-node sweep reverts the OTHER nodes' # progress; the caller scopes the failure to this node / re-queues its ids). if res is None or res.status_code != 200: - log.warning("[REMOVE] %-15s ==> non-200 %s: %s", nodes[node].name, - getattr(res, "status_code", "no-response"), getattr(res, "content", b"")[:200]) + log.warning( + "[REMOVE] %-15s ==> non-200 %s: %s", + nodes[node].name, + getattr(res, "status_code", "no-response"), + getattr(res, "content", b"")[:200], + ) return False # delete_many answers HTTP 200 even on a per-id failure; a GENUINE failure sets error=True / # status=partial_error (NOT an idempotent 'not exists', which stays error-absent). Surface that so @@ -384,8 +434,12 @@ def _delete_many(node, ids, nodes, db): except Exception: _body = {} if isinstance(_body, dict) and _body.get("error") is True: - log.warning("[REMOVE] %-15s ==> partial failure (%s): %s", nodes[node].name, - _body.get("status"), {k: v for k, v in _body.items() if v in ("error", "deleted_orphan_report")}) + log.warning( + "[REMOVE] %-15s ==> partial failure (%s): %s", + nodes[node].name, + _body.get("status"), + {k: v for k, v in _body.items() if v in ("error", "deleted_orphan_report")}, + ) return False return True @@ -676,7 +730,7 @@ def run(self): thread_targets.append((self.notification_loop, "notification_loop", ())) # Supervisor Loop - active_threads = {} # name -> thread_obj + active_threads = {} # name -> thread_obj log.info("Retriever supervisor started. Monitoring %d threads.", len(thread_targets)) @@ -752,12 +806,12 @@ def failed_cleaner(self): try: nodes = db.execute(select(Node).where(Node.enabled.is_(True))) for node in nodes or []: - failed_task_ids = [ - task["id"] - for task in 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] if not failed_task_ids: continue @@ -834,7 +888,16 @@ def fetcher(self): last_check = 0 last_checks[node.name] = 0 - task_ids = [task["id"] for task in node_fetch_tasks("reported", node.url, node.apikey, "fetch", last_check)] + tasks_data = node_fetch_tasks("reported", node.url, node.apikey, "fetch", last_check) + 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) + continue + else: + self.status_count[node.name] = 0 + + task_ids = [task["id"] for task in tasks_data] if task_ids: stmt = select(Task.task_id).where( @@ -854,17 +917,8 @@ def fetcher(self): queue_task_ids = set() for task_id in tasks_to_fetch: - try: - if task_id not in processed_task_ids and task_id not in queue_task_ids: - self.fetcher_queue.put(({"id": task_id}, node.id)) - except Exception as e: - self.status_count[node.name] += 1 - log.exception(e) - if self.status_count[node.name] == dead_count: - log.info("[-] %s dead", node.name) - # node_data = db.query(Node).filter_by(name=node.name).first() - # node_data.enabled = False - # db.commit() + if task_id not in processed_task_ids and task_id not in queue_task_ids: + self.fetcher_queue.put(({"id": task_id}, node.id)) db.commit() time.sleep(5) @@ -1070,7 +1124,6 @@ def fetch_latest_reports_nfs(self): self.current_queue[node_id].remove(task["id"]) db.commit() - def remove_from_worker(self): """ Removes tasks from worker nodes. @@ -1111,36 +1164,40 @@ def remove_from_worker(self): details.setdefault(node_id, set()).add(str(task_id)) if len(self.t_is_none.get(node_id)) > 50: break - db = session() - for node_id in details: - node = nodes[node_id] - if node and details.get(node_id): - ids = ",".join(list(set(details[node_id]))) - if not _delete_many(node_id, ids, nodes, db): - # The ids were consumed off the queue with .get() and this session has no writes to - # roll back, so a failed sweep would silently LOSE them (the worker's analyses/ leak). - # Re-queue so a frozen/down node's cleanup is retried on the next pass -- but BOUND it: - # retrying a stale (node_id, task_id) forever risks deleting a FRESH task after the - # worker is re-provisioned (node-local ids are reused). Drop (with a log) after - # CLEANER_MAX_RETRIES. Re-queue as int (details holds numeric strings) to match the - # int the t_is_none bookkeeping + original enqueue sites use. - for _tid in details[node_id]: - key = (node_id, int(_tid) if str(_tid).isdigit() else _tid) - attempts = self.cleaner_retries.get(key, 0) + 1 - if attempts >= CLEANER_MAX_RETRIES: - log.warning("[REMOVE] giving up on task %s @ node %s after %d failed cleanup " - "attempts (worker down/re-provisioned?)", _tid, node_id, attempts) - self.cleaner_retries.pop(key, None) - else: - self.cleaner_retries[key] = attempts - self.cleaner_queue.put(key) - else: - # Node cleanup succeeded -> clear any retry counters for its ids so the cap is per - # sustained-failure streak, not lifetime. - for _tid in details[node_id]: - self.cleaner_retries.pop((node_id, int(_tid) if str(_tid).isdigit() else _tid), None) - db.commit() - db.close() + with session() as db: + for node_id in details: + node = nodes[node_id] + if node and details.get(node_id): + ids = ",".join(list(set(details[node_id]))) + if not _delete_many(node_id, ids, nodes, db): + # The ids were consumed off the queue with .get() and this session has no writes to + # roll back, so a failed sweep would silently LOSE them (the worker's analyses/ leak). + # Re-queue so a frozen/down node's cleanup is retried on the next pass -- but BOUND it: + # retrying a stale (node_id, task_id) forever risks deleting a FRESH task after the + # worker is re-provisioned (node-local ids are reused). Drop (with a log) after + # CLEANER_MAX_RETRIES. Re-queue as int (details holds numeric strings) to match the + # int the t_is_none bookkeeping + original enqueue sites use. + for _tid in details[node_id]: + key = (node_id, int(_tid) if str(_tid).isdigit() else _tid) + attempts = self.cleaner_retries.get(key, 0) + 1 + if attempts >= CLEANER_MAX_RETRIES: + log.warning( + "[REMOVE] giving up on task %s @ node %s after %d failed cleanup " + "attempts (worker down/re-provisioned?)", + _tid, + node_id, + attempts, + ) + self.cleaner_retries.pop(key, None) + else: + self.cleaner_retries[key] = attempts + self.cleaner_queue.put(key) + else: + # Node cleanup succeeded -> clear any retry counters for its ids so the cap is per + # sustained-failure streak, not lifetime. + for _tid in details[node_id]: + self.cleaner_retries.pop((node_id, int(_tid) if str(_tid).isdigit() else _tid), None) + db.commit() time.sleep(20) @@ -1238,7 +1295,11 @@ def submit_tasks(self, node_name, pend_tasks_num, options_like=False, force_push bin_path = os.path.join(CUCKOO_ROOT, "storage", "binaries", sample_sha256) if sample_sha256 else None if bin_path and path_exists(bin_path): - log.info("Task id: %d - Target file not found at original path, but found in binaries storage: %s. Updating target path.", t.id, bin_path) + log.info( + "Task id: %d - Target file not found at original path, but found in binaries storage: %s. Updating target path.", + t.id, + bin_path, + ) t.target = bin_path else: log.info("Task id: %d - File doesn't exist: %s", t.id, t.target) @@ -1248,9 +1309,7 @@ def submit_tasks(self, node_name, pend_tasks_num, options_like=False, force_push # We can't upload size bigger than X to our workers. In case we extract archive that contains bigger file. file_size = path_get_size(t.target) if file_size > web_conf.general.max_sample_size: - log.debug( - "File size: %d is bigger than allowed: %d", file_size, web_conf.general.max_sample_size - ) + log.debug("File size: %d is bigger than allowed: %d", file_size, web_conf.general.max_sample_size) main_db.set_status(t.id, TASK_BANNED) continue options = get_options(t.options) @@ -1360,56 +1419,26 @@ def submit_tasks(self, node_name, pend_tasks_num, options_like=False, force_push db.commit() log.info("Pushed all tasks") return True - # ToDo not finished - # Only get tasks that have not been pushed yet. - """ - q = db.query(Task).filter(or_(Task.node_id.is_(None), Task.task_id.is_(None)), Task.finished.is_(False)) - if q is None: - db.commit() - return True - # Order by task priority and task id. - q = q.order_by(-Task.priority, Task.main_task_id) - # if we have node set in options push - if dist_conf.distributed.enable_tags: - # Create filter query from tasks in ta - tags = [getattr(Task, "tags") == ""] - for tg in SERVER_TAGS[node.name]: - if len(tg.split(",")) == 1: - tags.append(getattr(Task, "tags") == (tg + ",")) - else: - tg = tg.split(",") - # ie. LIKE "%,%,%," - t_combined = [getattr(Task, "tags").like("%s" % ("%," * len(tg)))] - for tag in tg: - t_combined.append(getattr(Task, "tags").like("%%%s%%" % (tag + ","))) - tags.append(and_(*t_combined)) - # Filter by available tags - q = q.filter(or_(*tags)) - to_upload = q.limit(pend_tasks_num).all() - """ # 1. Start with a select() statement and initial filters. stmt = ( select(Task) .where(or_(Task.node_id.is_(None), Task.task_id.is_(None)), Task.finished.is_(False)) .order_by(Task.priority.desc(), Task.main_task_id) ) - # print(stmt, "stmt") - # ToDo broken - """ - # 3. Apply the dynamic tag filter. - if dist_conf.distributed.enable_tags: + + # 2. Apply the dynamic tag filter. + if dist_conf.distributed.enable_tags and node.name in SERVER_TAGS: tags_conditions = [Task.tags == ""] for tg in SERVER_TAGS[node.name]: tags_list = tg.split(",") if len(tags_list) == 1: tags_conditions.append(Task.tags == f"{tg},") else: - # The pattern of building a list of conditions for `and_` or `or_` - # works the same way with the modern .where() clause. - t_combined = [Task.tags.like(f"%{tag},%") for tag in tags_list] + t_combined = [Task.tags.like("%" + "%," * len(tags_list))] + for tag in tags_list: + t_combined.append(Task.tags.like(f"%{tag},%")) tags_conditions.append(and_(*t_combined)) stmt = stmt.where(or_(*tags_conditions)) - """ # 4. Apply the limit and execute the query. to_upload = db.scalars(stmt.limit(pend_tasks_num)).all() # print(to_upload, node.name, pend_tasks_num) @@ -1421,8 +1450,7 @@ def submit_tasks(self, node_name, pend_tasks_num, options_like=False, force_push max_workers = int(dist_conf.distributed.dist_threads) with ThreadPoolExecutor(max_workers=max_workers) as executor: future_to_task = { - executor.submit(node_submit_task, task.id, node.id, task.main_task_id, db=None): task - for task in to_upload + executor.submit(node_submit_task, task.id, node.id, task.main_task_id, db=None): task for task in to_upload } count_submitted = 0 for future in as_completed(future_to_task): @@ -1557,9 +1585,11 @@ def run(self): # Balance the tasks, works fine if no tags are set node_name = min( STATUSES, - key=lambda k: STATUSES[k]["tasks"]["completed"] - + STATUSES[k]["tasks"]["pending"] - + STATUSES[k]["tasks"]["running"], + key=lambda k: ( + STATUSES[k]["tasks"]["completed"] + + STATUSES[k]["tasks"]["pending"] + + STATUSES[k]["tasks"]["running"] + ), ) if node_name != node.name: node = db.scalar(select(Node).where(Node.name == node_name)) @@ -1595,255 +1625,53 @@ def run(self): time.sleep(INTERVAL) -def output_json(data, code, headers=None): - """ - Create a JSON response with the given data, HTTP status code, and optional headers. +# FastAPI endpoints will be registered dynamically inside create_app() below - Args: - data (dict): The data to be serialized to JSON. - code (int): The HTTP status code for the response. - headers (dict, optional): Additional headers to include in the response. Defaults to None. - - Returns: - Response: A Flask response object with the JSON data and specified headers. - """ - resp = make_response(json.dumps(data), code) - resp.headers.extend(headers or {}) - return resp +def update_machine_table(node_name): + with session() as db: + node = db.scalar(select(Node).where(Node.name == node_name)) -class NodeBaseApi(RestResource): - def __init__(self, *args, **kwargs): - RestResource.__init__(self, *args, **kwargs) - - self._parser = reqparse.RequestParser() - self._parser.add_argument("name", type=str, location="form") - self._parser.add_argument("url", type=str, location="form") - self._parser.add_argument("apikey", type=str, default="", location="form") - self._parser.add_argument("exitnodes", type=distutils.util.strtobool, default=None, location="form") - self._parser.add_argument("enabled", type=distutils.util.strtobool, default=None, location="form") + # get new vms + new_machines = node_list_machines(node.url, node.apikey) + # delete all old vms + db.execute(delete(Machine).where(Machine.node_id == node.id)) -class NodeRootApi(NodeBaseApi): - def get(self): - nodes = {} - db = session() - for node in db.scalars(select(Node)): - machines = [ - dict( - name=machine.name, - platform=machine.platform, - tags=machine.tags, - ) - for machine in node.machines.all() - ] - - nodes[node.name] = dict( - name=node.name, - url=node.url, - machines=machines, - enabled=node.enabled, - ) - db.close() - return dict(nodes=nodes) - - def post(self): - db = session() - args = self._parser.parse_args() - node_exist = False - # On autoscaling we might get the same name but different IP for server. Kinda PUT friendly POST - node = db.scalar(select(Node).where(Node.name == args["name"])) - if node: - if node.url == args["url"]: - return dict(success=False, message=f"Node called {args['name']} already exists") - else: - node.url = args["url"] - else: - node = Node(name=args["name"], url=args["url"], apikey=args["apikey"]) - - machines = [] - for machine in node_list_machines(args["url"], args["apikey"]): - machines.append(dict(name=machine.name, platform=machine.platform, tags=machine.tags)) + log.info("Available VM's on %s:", node_name) + # replace with new vms + for machine in new_machines: + log.info("-->\t%s", machine.name) node.machines.append(machine) db.add(machine) - exitnodes = [] - for exitnode in node_list_exitnodes(args["url"], args.get("apikey")): - exitnode_db = db.scalar(select(ExitNodes).where(ExitNodes.name == exitnode)) - if exitnode_db: - exitnode = exitnode_db - else: - exitnode = ExitNodes(name=exitnode) - exitnodes.append(dict(name=exitnode.name)) - node.exitnodes.append(exitnode) - db.add(exitnode) - - if args.get("enabled"): - node.enabled = bool(args["enabled"]) - - if not node_exist: - db.add(node) - db.commit() - db.close() - - if NFS_FETCH: - # Add entry to /etc/fstab, create folder and mount server - hostname = urlparse(args["url"]).netloc.split(":")[0] - if hostname != main_server_name: - send_socket_command(dist_conf.NFS.fstab_socket, "add_entry", *[hostname, args["name"]]) - - return dict(name=args["name"], machines=machines, exitnodes=exitnodes) - - -class NodeApi(NodeBaseApi): - def get(self, name): - db = session() - node = db.scalar(select(Node).where(Node.name == name)) - db.close() - return dict(name=node.name, url=node.url) - - def put(self, name): - db = session() - args = self._parser.parse_args() - node = db.scalar(select(Node).where(Node.name == name)) - - if not node: - return dict(error=True, error_value="Node doesn't exist") - - for k, v in args.items(): - if k == "exitnodes": - exitnodes = [] - for exitnode in node_list_exitnodes(node.url, node.apikey): - exitnode_db = db.scalar(select(ExitNodes).where(ExitNodes.name == exitnode)) - if exitnode_db: - exitnode = exitnode_db - else: - exitnode = ExitNodes(name=exitnode) - exitnodes.append(dict(name=exitnode.name)) - node.exitnodes.append(exitnode) - db.add(exitnode) - db.add(node) - else: - if v is not None: - setattr(node, k, v) - db.commit() - db.close() - return dict(error=False, error_value=f"Successfully modified node: {name}") - - def delete(self, name): - db = session() - node = db.scalar(select(Node).where(Node.name == name)) - node.enabled = False db.commit() - db.close() - - -class TaskBaseApi(RestResource): - def __init__(self, *args, **kwargs): - RestResource.__init__(self, *args, **kwargs) - - self._parser = reqparse.RequestParser() - self._parser.add_argument("package", type=str, default="", location="form") - self._parser.add_argument("timeout", type=int, default=0, location="form") - self._parser.add_argument("priority", type=int, default=1, location="form") - self._parser.add_argument("options", type=str, default="", location="form") - self._parser.add_argument("machine", type=str, default="", location="form") - self._parser.add_argument("platform", type=str, default="windows", location="form") - self._parser.add_argument("tags", type=str, default="", location="form") - self._parser.add_argument("custom", type=str, default="", location="form") - self._parser.add_argument("memory", type=str, default="0", location="form") - self._parser.add_argument("clock", type=int, location="form") - self._parser.add_argument("enforce_timeout", type=bool, default=False, location="form") - - -class TaskInfo(RestResource): - def get(self, main_task_id): - response = {"status": 0} - db = session() - task_db = db.scalar(select(Task).where(Task.main_task_id == main_task_id)) - if task_db and task_db.node_id: - node_stmt = select(Node).where(Node.id == task_db.node_id) - node = db.scalar(node_stmt) - response = {"status": 1, "task_id": task_db.task_id, "url": node.url, "name": node.name} - else: - response = {"status": "pending"} - db.close() - return response - - -class StatusRootApi(RestResource): - def get(self): - # null = None - db = session() - unified_counts = db.execute( - select( - func.count(case((and_(Task.node_id.is_not(None), Task.finished.is_(False)), Task.id))).label("processing"), - func.count(case((and_(Task.node_id.is_not(None), Task.finished.is_(True)), Task.id))).label("processed"), - func.count(case((Task.node_id.is_(None), Task.id))).label("pending"), - ) - ).first() - tasks_counts = { - "processing": unified_counts.processing, - "processed": unified_counts.processed, - "pending": unified_counts.pending, - } - return jsonify({"nodes": STATUSES, "tasks": tasks_counts}) - - -class DistRestApi(RestApi): - def __init__(self, *args, **kwargs): - RestApi.__init__(self, *args, **kwargs) - self.representations = { - "application/json": output_json, - } - - -def update_machine_table(node_name): - db = session() - node = db.scalar(select(Node).where(Node.name == node_name)) - - # get new vms - new_machines = node_list_machines(node.url, node.apikey) - - # delete all old vms - db.execute(delete(Machine).where(Machine.node_id == node.id)) - - log.info("Available VM's on %s:", node_name) - # replace with new vms - for machine in new_machines: - log.info("-->\t%s", machine.name) - node.machines.append(machine) - db.add(machine) - - db.commit() log.info("Updated the machine table for node: %s", node_name) def delete_vm_on_node(node_name, vm_name): - db = session() - node = db.scalar(select(Node).where(Node.name == node_name)) - vm = db.scalar(select(Machine).where(Machine.name == vm_name, Machine.node_id == node.id)) + with session() as db: + node = db.scalar(select(Node).where(Node.name == node_name)) + vm = db.scalar(select(Machine).where(Machine.name == vm_name, Machine.node_id == node.id)) - if not vm: - log.error("The selected VM does not exist") - return + if not vm: + log.error("The selected VM does not exist") + return - status = node.delete_machine(vm_name) + status = node.delete_machine(vm_name) - if status: - # delete vm in dist db - db.execute(delete(Machine).where(Machine.name == vm_name, Machine.node_id == node.id)) - db.commit() - db.close() + if status: + # delete vm in dist db + db.execute(delete(Machine).where(Machine.name == vm_name, Machine.node_id == node.id)) + db.commit() def node_enabled(node_name, status): - db = session() - node = db.scalar(select(Node).where(Node.name == node_name)) - node.enabled = status - db.commit() - db.close() + with session() as db: + node = db.scalar(select(Node).where(Node.name == node_name)) + node.enabled = status + db.commit() def cron_cleaner(clean_x_hours=False): @@ -1983,18 +1811,167 @@ def cron_cleaner(clean_x_hours=False): def create_app(database_connection): - # http://flask-sqlalchemy.pocoo.org/2.1/config/ - # https://github.com/tmeryu/flask-sqlalchemy/blob/master/flask_sqlalchemy/__init__.py#L787 - app = Flask("Distributed CAPE") - # app.config["SQLALCHEMY_DATABASE_URI"] = database_connection - app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = True - app.config["SQLALCHEMY_POOL_SIZE"] = int(dist_conf.distributed.dist_threads) + 5 - app.config["SECRET_KEY"] = os.urandom(32) - restapi = DistRestApi(app) - restapi.add_resource(NodeRootApi, "/node") - restapi.add_resource(NodeApi, "/node/") - restapi.add_resource(StatusRootApi, "/status") - restapi.add_resource(TaskInfo, "/task/") + from pydantic import BaseModel + + class NodeRegister(BaseModel): + name: str + url: str + apikey: str = "" + enabled: Optional[bool] = None + + class NodeUpdate(BaseModel): + url: Optional[str] = None + apikey: Optional[str] = None + exitnodes: Optional[bool] = None + enabled: Optional[bool] = None + + app = FastAPI(title="Distributed CAPE") + + @app.get("/node") + def get_nodes(): + nodes = {} + with session() as db: + for node in db.scalars(select(Node)): + machines = [ + dict( + name=machine.name, + platform=machine.platform, + tags=machine.tags, + ) + for machine in node.machines.all() + ] + + nodes[node.name] = dict( + name=node.name, + url=node.url, + machines=machines, + enabled=node.enabled, + ) + return dict(nodes=nodes) + + @app.post("/node") + def post_node(payload: NodeRegister): + with session() as db: + node_exist = False + # On autoscaling we might get the same name but different IP for server. Kinda PUT friendly POST + node = db.scalar(select(Node).where(Node.name == payload.name)) + if node: + node_exist = True + if node.url == payload.url: + return dict(success=False, message=f"Node called {payload.name} already exists") + else: + node.url = payload.url + else: + node = Node(name=payload.name, url=payload.url, apikey=payload.apikey) + + machines = [] + for machine in node_list_machines(payload.url, payload.apikey): + machines.append(dict(name=machine.name, platform=machine.platform, tags=machine.tags)) + node.machines.append(machine) + db.add(machine) + + exitnodes = [] + for exitnode in node_list_exitnodes(payload.url, payload.apikey): + exitnode_db = db.scalar(select(ExitNodes).where(ExitNodes.name == exitnode)) + if exitnode_db: + exitnode = exitnode_db + else: + exitnode = ExitNodes(name=exitnode) + exitnodes.append(dict(name=exitnode.name)) + node.exitnodes.append(exitnode) + db.add(exitnode) + + if payload.enabled is not None: + node.enabled = payload.enabled + + if not node_exist: + db.add(node) + db.commit() + + if NFS_FETCH: + # Add entry to /etc/fstab, create folder and mount server + hostname = urlparse(payload.url).netloc.split(":")[0] + if hostname != main_server_name: + send_socket_command(dist_conf.NFS.fstab_socket, "add_entry", *[hostname, payload.name]) + + return dict(name=payload.name, machines=machines, exitnodes=exitnodes) + + @app.get("/node/{name}") + def get_node(name: str): + with session() as db: + node = db.scalar(select(Node).where(Node.name == name)) + if not node: + raise HTTPException(status_code=404, detail="Node doesn't exist") + return dict(name=node.name, url=node.url) + + @app.put("/node/{name}") + def put_node(name: str, payload: NodeUpdate): + with session() as db: + node = db.scalar(select(Node).where(Node.name == name)) + + if not node: + return dict(error=True, error_value="Node doesn't exist") + + if payload.url is not None: + node.url = payload.url + if payload.apikey is not None: + node.apikey = payload.apikey + if payload.enabled is not None: + node.enabled = payload.enabled + + if payload.exitnodes: + exit_list = [] + for exitnode in node_list_exitnodes(node.url, node.apikey): + exitnode_db = db.scalar(select(ExitNodes).where(ExitNodes.name == exitnode)) + if exitnode_db: + exitnode = exitnode_db + else: + exitnode = ExitNodes(name=exitnode) + exit_list.append(dict(name=exitnode.name)) + node.exitnodes.append(exitnode) + db.add(exitnode) + db.add(node) + db.commit() + return dict(error=False, error_value=f"Successfully modified node: {name}") + + @app.delete("/node/{name}") + def delete_node(name: str): + with session() as db: + node = db.scalar(select(Node).where(Node.name == name)) + if node: + node.enabled = False + db.commit() + return None + + @app.get("/status") + def get_status(): + with session() as db: + unified_counts = db.execute( + select( + func.count(case((and_(Task.node_id.is_not(None), Task.finished.is_(False)), Task.id))).label("processing"), + func.count(case((and_(Task.node_id.is_not(None), Task.finished.is_(True)), Task.id))).label("processed"), + func.count(case((Task.node_id.is_(None), Task.id))).label("pending"), + ) + ).first() + tasks_counts = { + "processing": unified_counts.processing, + "processed": unified_counts.processed, + "pending": unified_counts.pending, + } + return {"nodes": STATUSES, "tasks": tasks_counts} + + @app.get("/task/{main_task_id}") + def get_task_info(main_task_id: int): + response = {"status": 0} + with session() as db: + task_db = db.scalar(select(Task).where(Task.main_task_id == main_task_id)) + if task_db and task_db.node_id: + node_stmt = select(Node).where(Node.id == task_db.node_id) + node = db.scalar(node_stmt) + response = {"status": 1, "task_id": task_db.task_id, "url": node.url, "name": node.name} + else: + response = {"status": "pending"} + return response return app @@ -2026,78 +2003,160 @@ def init_logging(debug=False): return log -if __name__ == "__main__": - p = argparse.ArgumentParser() - p.add_argument("host", nargs="?", default="0.0.0.0", help="Host to listen on") - p.add_argument("port", nargs="?", type=int, default=9003, help="Port to listen on") - p.add_argument("-d", "--debug", action="store_true", help="Enable debug logging") - p.add_argument("--uptime-logfile", type=str, help="Uptime logfile path") - p.add_argument("--node", type=str, help="Node name to update in distributed DB") - p.add_argument("--delete-vm", type=str, help="VM name to delete from Node") - p.add_argument("--disable", action="store_true", help="Disable Node provided in --node") - p.add_argument("--enable", action="store_true", help="Enable Node provided in --node") - p.add_argument("--clean-workers", action="store_true", help="Delete reported and notificated tasks from workers") - p.add_argument( - "-ec", - "--enable-clean", - action="store_true", - help="Enable delete tasks from nodes, also will remove tasks submited by humands and not dist", - ) - p.add_argument( - "-ef", - "--enable-failed-clean", - action="store_true", - default=False, - help="Enable delete failed tasks from nodes, also will remove tasks submited by humands and not dist", - ) - p.add_argument("-fr", "--force-reported", action="store", help="change report to reported") - p.add_argument( - "-ch", - "--clean-hours", - action="store", - type=int, - default=0, - help="Clean tasks for last X hours", - ) - p.add_argument( - "--submit-only", - action="store_true", - help="Disable retrieval threads (use when running Go Fast-Fetcher)", - ) - p.add_argument( - "--gcs-replay", - action="store", - help="Replay GCS upload for a range of tasks (e.g., 1-100 or 1,2,3)", - ) - p.add_argument( - "--gcs-sync", - action="store", - help="Sync GCS with DB for a given time range (e.g., 12h, 1d, 2d)", - ) - p.add_argument( - "--gcs-refetch-banned", - action="store", - help="Refetch banned tasks from GCS for a given time range (e.g., 12h, 1d, 2d)", +def show_status_cli(): + with session() as db: + unified_counts = db.execute( + select( + func.count(case((and_(Task.node_id.is_not(None), Task.finished.is_(False)), Task.id))).label("processing"), + func.count(case((and_(Task.node_id.is_not(None), Task.finished.is_(True)), Task.id))).label("processed"), + func.count(case((Task.node_id.is_(None), Task.id))).label("pending"), + ) + ).first() + print("Cluster Status Summary:") + print("-" * 30) + print(f"Processing tasks : {unified_counts.processing}") + print(f"Processed tasks : {unified_counts.processed}") + print(f"Pending tasks : {unified_counts.pending}") + print("-" * 30) + print("Active Nodes:") + for node_name, status in STATUSES.items(): + print(f" Node: {node_name:<15} | Status: {status.get('status', 'unknown')}") + + +def list_nodes_cli(): + with session() as db: + nodes = db.scalars(select(Node)).all() + if not nodes: + print("No registered CAPE nodes found.") + return + print(f"{'Node Name':<20} | {'Enabled':<8} | {'API URL':<50}") + print("-" * 84) + for node in nodes: + print(f"{node.name:<20} | {str(node.enabled):<8} | {node.url:<50}") + machines = node.machines.all() + if machines: + print(" Virtual Machines:") + for machine in machines: + print(f" --> {machine.name} ({machine.platform}) - Tags: {machine.tags}") + print("-" * 84) + + +def register_node_cli(name, url, apikey, enabled): + if not name or not url: + print("Error: Registering a node requires both --node and --url ") + sys.exit(1) + with session() as db: + node_exist = False + node = db.scalar(select(Node).where(Node.name == name)) + if node: + node_exist = True + if node.url == url: + print(f"Node called {name} already exists.") + return + else: + node.url = url + else: + node = Node(name=name, url=url, apikey=apikey) + + machines = [] + try: + for machine in node_list_machines(url, apikey): + machines.append(machine) + node.machines.append(machine) + db.add(machine) + except Exception as e: + print(f"Warning: Could not fetch machines from node: {e}") + + try: + for exitnode in node_list_exitnodes(url, apikey): + exitnode_db = db.scalar(select(ExitNodes).where(ExitNodes.name == exitnode)) + if exitnode_db: + exitnode = exitnode_db + else: + exitnode = ExitNodes(name=exitnode) + node.exitnodes.append(exitnode) + db.add(exitnode) + except Exception as e: + print(f"Warning: Could not fetch exitnodes from node: {e}") + + if enabled is not None: + node.enabled = enabled + + if not node_exist: + db.add(node) + db.commit() + print(f"Successfully registered node '{name}' with {len(machines)} machines.") + + +def modify_node_cli(name, url=None, apikey=None, enabled=None): + if not name: + print("Error: Modifying a node requires --node ") + sys.exit(1) + with session() as db: + node = db.scalar(select(Node).where(Node.name == name)) + if not node: + print(f"Error: Node '{name}' does not exist.") + sys.exit(1) + if url is not None: + node.url = url + if apikey is not None: + node.apikey = apikey + if enabled is not None: + node.enabled = enabled + db.commit() + print(f"Successfully modified node '{name}'.") + + +def main(): + p = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, description="Distributed CAPE Daemon & Admin Utility" ) - p.add_argument( - "--samples-bucket", - action="store", - help="Specify GCS bucket for samples (used with --gcs-refetch-banned)", + + g_daemon = p.add_argument_group("Daemon / Server Options") + g_daemon.add_argument("host", nargs="?", default="0.0.0.0", help="Host to listen on") + g_daemon.add_argument("port", nargs="?", type=int, default=9003, help="Port to listen on") + g_daemon.add_argument("-d", "--debug", action="store_true", help="Enable debug logging") + g_daemon.add_argument("--uptime-logfile", type=str, help="Uptime logfile path") + g_daemon.add_argument("--submit-only", action="store_true", help="Disable retrieval threads") + g_daemon.add_argument("-ec", "--enable-clean", action="store_true", help="Enable delete tasks from nodes") + g_daemon.add_argument( + "-ef", "--enable-failed-clean", action="store_true", default=False, help="Enable delete failed tasks from nodes" ) + g_daemon.add_argument("-ch", "--clean-hours", action="store", type=int, default=0, help="Clean tasks for last X hours") + + g_admin = p.add_argument_group("Cluster Administration (CLI Admin Tools)") + g_admin.add_argument("--status", action="store_true", help="Show cluster status summary") + g_admin.add_argument("--list-nodes", action="store_true", help="List all registered nodes and their VMs") + g_admin.add_argument("--register-node", action="store_true", help="Register a new CAPE worker node (requires --node, --url)") + g_admin.add_argument("--modify-node", action="store_true", help="Modify an existing registered node (requires --node)") + + g_vm = p.add_argument_group("VM Maintenance & Node Control") + g_vm.add_argument("--node", type=str, help="Node name to update, register, or modify") + g_vm.add_argument("--url", type=str, help="API URL of the node (for register / modify)") + g_vm.add_argument("--apikey", type=str, default="", help="API Key of the node (for register / modify)") + g_vm.add_argument("--delete-vm", type=str, help="VM name to delete from Node") + g_vm.add_argument("--disable", action="store_true", help="Disable/deactivate the Node") + g_vm.add_argument("--enable", action="store_true", help="Enable/activate the Node") + + g_sync = p.add_argument_group("Storage, GCS, & Sync Options") + g_sync.add_argument("--clean-workers", action="store_true", help="Delete reported tasks from workers") + g_sync.add_argument("-fr", "--force-reported", action="store", help="Change task status to reported") + g_sync.add_argument("--gcs-replay", action="store", help="Replay GCS upload for a range of tasks") + g_sync.add_argument("--gcs-sync", action="store", help="Sync GCS with DB for a given time range") + g_sync.add_argument("--gcs-refetch-banned", action="store", help="Refetch banned tasks from GCS") + g_sync.add_argument("--samples-bucket", action="store", help="Specify GCS bucket for samples") args = p.parse_args() + global log log = init_logging(args.debug) init_database() if args.enable_clean: cron_cleaner(args.clean_hours) - # sys.exit() if args.force_reported: with main_db.session.begin(): - # set completed_on time main_db.set_status(args.force_reported, TASK_DISTRIBUTED_COMPLETED) - # set reported time main_db.set_status(args.force_reported, TASK_REPORTED) sys.exit() @@ -2113,8 +2172,29 @@ def init_logging(debug=False): gcs_refetch_banned(args.gcs_refetch_banned, samples_bucket=args.samples_bucket) sys.exit() + # CLI Admin Commands execution + if args.status: + show_status_cli() + sys.exit() + + if args.list_nodes: + list_nodes_cli() + sys.exit() + + if args.register_node: + enabled_val = True if args.enable else (False if args.disable else True) + register_node_cli(args.node, args.url, args.apikey, enabled_val) + sys.exit() + + if args.modify_node: + enabled_val = True if args.enable else (False if args.disable else None) + modify_node_cli(args.node, args.url, args.apikey, enabled_val) + sys.exit() + + global delete_enabled, failed_clean_enabled delete_enabled = args.enable_clean failed_clean_enabled = args.enable_failed_clean + if args.node: if args.delete_vm: delete_vm_on_node(args.node, args.delete_vm) @@ -2126,22 +2206,27 @@ def init_logging(debug=False): update_machine_table(args.node) sys.exit() + # Starts Daemon Server + app = create_app(database_connection=dist_conf.distributed.db) + + t = StatusThread(name="StatusThread") + t.daemon = True + t.start() + + if not args.submit_only and not dist_conf.distributed.get("submit_only"): + retrieve = Retriever(name="Retriever") + retrieve.daemon = True + retrieve.start() else: - app = create_app(database_connection=dist_conf.distributed.db) + log.info("Submit-only mode: Retriever thread disabled.") - t = StatusThread(name="StatusThread") - t.daemon = True - t.start() + import uvicorn - if not args.submit_only and not dist_conf.distributed.get("submit_only"): - retrieve = Retriever(name="Retriever") - retrieve.daemon = True - retrieve.start() - else: - log.info("Submit-only mode: Retriever thread disabled.") + uvicorn.run(app, host=args.host, port=args.port, log_level="debug" if args.debug else "info") - app.run(host=args.host, port=args.port, debug=args.debug, use_reloader=False) +if __name__ == "__main__": + main() else: init_database(exists_ok=True) app = create_app(database_connection=dist_conf.distributed.db)