clue-api 1.3.0.dev92__tar.gz → 1.3.0.dev102__tar.gz

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: clue-api
-Version: 1.3.0.dev92
+Version: 1.3.0.dev102
 Summary: Clue distributed enrichment service
 License: MIT
 License-File: LICENSE
@@ -24,6 +24,7 @@ Requires-Dist: authlib (<2.0.0) ; extra == "server"
 Requires-Dist: bcrypt (>=4.1.2,<5.0.0) ; extra == "server"
 Requires-Dist: beautifulsoup4 (>=4.13.3,<5.0.0)
 Requires-Dist: cart (>=1.2.3,<2.0.0)
+Requires-Dist: celery (>=5.6.2,<6.0.0)
 Requires-Dist: elastic-apm (>=6.22.0,<7.0.0)
 Requires-Dist: flasgger (>=0.9.7.1,<0.10.0.0) ; extra == "server"
 Requires-Dist: flask (<3.0.0)

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/api/v1/actions.py

@@ -90,3 +90,36 @@ def execute_action(plugin_id: str, action_id: str, **kwargs) -> ActionResult:
         return not_found(err=err.message)
     except ClueException as err:
         return internal_error(err=err.message)
+
+
+@generate_swagger_docs(responses={200: "Successfully fetched status of action"})
+@actions_api.route("/<plugin_id>/<action_id>/status/<task_id>", methods=["GET"])
+@api_login()
+def get_action_status(plugin_id: str, action_id: str, task_id: str, **kwargs) -> ActionResult:
+    """Get the status or result of a running action.
+
+    Variables:
+    plugin_id (str): the ID of the plugin who owns the action to execute
+    action_id (str): the ID of the action to execute
+    task_id (str): the ID of the specific task to get the status of
+
+    Arguments:
+    task_id (str): the celery task id to get the status of
+
+
+    Result Example:
+    {
+        "outcome": "success | failure | pending", # was this execution a success or failure or is it still pending?
+        "format": "link", # What format is the output in?
+        "output": "http://example.com" # The output of the action. Can be any data structure.
+        "task_id": if the action is still running, what is the task id so that we can fetch the status again
+    }
+    """
+    try:
+        if not task_id:
+            return internal_error(err="no task_id found in url. task_id is required for this request.")
+        return ok(action_service.get_action_status(plugin_id, action_id, task_id, kwargs["user"]))
+    except NotFoundException as err:
+        return not_found(err=err.message)
+    except ClueException as err:
+        return internal_error(err=err.message)

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/app.py

@@ -2,6 +2,8 @@ import warnings
 
 from gevent import monkey
 
+from clue.constants.supported_types import SUPPORTED_TYPES
+
 monkey.patch_all()
 
 import os
@@ -132,7 +134,8 @@ app.register_blueprint(registration_api)
 app.register_blueprint(static_api)
 
 
-logger.info("Checking extensions for additional routes")
+logger.info("Checking extensions for initialization and additional routes")
+num_buildin_types = len(SUPPORTED_TYPES)
 for extension in get_extensions():
     if extension.modules.init:
         extension.modules.init(flask_app=app)
@@ -144,6 +147,8 @@ for extension in get_extensions():
         logger.info("Enabling additional endpoint: %s", route.url_prefix)
         app.register_blueprint(route)
 
+logger.info("%s types configured (%s custom types)", len(SUPPORTED_TYPES), len(SUPPORTED_TYPES) - num_buildin_types)
+
 # Setup OAuth providers
 if config.auth.oauth.enabled:
     providers = []

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/common/regex.py

@@ -40,5 +40,3 @@ URI_ONLY = f"^{URI_REGEX}$"
 UUID4_REGEX = r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-4[0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$"
 
 EMAIL_PATH_REGEX = r"^[A-Z]+_EMAIL://.*"
-HBS_AGENT_ID_REGEX = r"[0-9a-fA-F]{1,4}\.[0-9a-fA-F]{1,4}\.[0-9a-fA-F]{1,4}\.[0-9a-fA-F]{1,4}"
-HBS_AGENT_ID_ONLY_REGEX = f"^{HBS_AGENT_ID_REGEX}$"

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/constants/supported_types.py

@@ -3,7 +3,6 @@ from clue.common.regex import (
     DOMAIN_ONLY_REGEX,
     EMAIL_PATH_REGEX,
     EMAIL_REGEX,
-    HBS_AGENT_ID_REGEX,
     IPV4_ONLY_REGEX,
     IPV6_ONLY_REGEX,
     MD5_REGEX,
@@ -32,16 +31,17 @@ SUPPORTED_TYPES = {
     "md5": MD5_REGEX,
     "sha1": SHA1_REGEX,
     "sha256": SHA256_REGEX,
-    "hbs_oid": None,
-    "hbs_agent_id": HBS_AGENT_ID_REGEX,
     "telemetry": None,
-    "howler_id": None,
     "hostname": None,
     "tenant-id": UUID4_REGEX,
 }
 
+CASE_INSENSITIVE_TYPES = ["ip", "domain", "port", "tenant-id", "hbs_oid", "hbs_agent_id"]
+
 
-def add_supported_type(type: str, regex: str | None = None, namespace: str | None = None):
+def add_supported_type(
+    type: str, regex: str | None = None, namespace: str | None = None, case_insensitive: bool = False
+):
     r"""Add a supported type to the SUPPORTED_TYPES registry.
 
     This function registers a new type with an optional regex pattern for validation.
@@ -62,10 +62,11 @@ def add_supported_type(type: str, regex: str | None = None, namespace: str | Non
     """
     if not namespace:
         logger.info("Adding new type %s to the default namespace with regex %s", type, regex)
-        SUPPORTED_TYPES[type] = regex
+        new_entry = type
     else:
         logger.info("Adding type %s to namespace %s with regex %s", type, namespace, regex)
-        SUPPORTED_TYPES[f"{namespace}/{type}"] = regex
+        new_entry = f"{namespace}/{type}"
 
-
-CASE_INSENSITIVE_TYPES = ["ip", "domain", "port", "tenant-id", "hbs_oid", "hbs_agent_id"]
+    SUPPORTED_TYPES[new_entry] = regex
+    if case_insensitive:
+        CASE_INSENSITIVE_TYPES.append(new_entry)

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/extensions/config.py

@@ -59,6 +59,10 @@ class BaseExtensionConfig(BaseSettings):
 
             data["modules"]["routes"] = new_routes
 
+        if "init" in data["modules"]:
+            if isinstance(data["modules"]["init"], bool):
+                data["modules"]["init"] = f"{plugin_name}.init:initialize"
+
         if "obo_module" in data["modules"]:
             if isinstance(data["modules"]["obo_module"], bool):
                 data["modules"]["obo_module"] = f"{plugin_name}.obo:get_obo_token"

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/models/actions.py

@@ -69,6 +69,10 @@ class ActionContextInformation(BaseModel):
 ActionContextInformationType = TypeVar("ActionContextInformationType", bound=ActionContextInformation)
 
 
+class ActionStatusRequest(BaseModel):
+    task_id: str = Field(description="The task id to get the status for.")
+
+
 class ExecuteRequest(BaseModel):
     context: ActionContextInformation | None = Field(
         description="Contextual information on where the action is being executed (if provided)", default=None
@@ -123,6 +127,7 @@ class ActionBase(BaseModel):
     )
     accept_empty: bool = Field(description="Does this action support execution with no selectors?", default=False)
     accept_multiple: bool = Field(description="Does this action support multiple values?", default=False)
+    async_result: bool = Field(description="Does this action run asynchronously?", default=False)
     format: str | None = Field(
         description="What is the format of the output, if known?",
         default=None,
@@ -242,7 +247,9 @@ class Action(ActionBase, Generic[ER]):
 
 
 class ActionResult(BaseModel, Generic[DATA]):
-    outcome: Union[Literal["success"], Literal["failure"]] = Field(description="Did the action succeed or fail?")
+    outcome: Union[Literal["success"], Literal["failure"], Literal["pending"]] = Field(
+        description="Did the action succeed/fail, or is it pending?"
+    )
     summary: str | None = Field(description="Message explaining the outcome of the action.", default=None)
     output: DATA | Url | None = Field(description="The output of the action.", default=None)
     format: str | None = Field(
@@ -251,6 +258,7 @@ class ActionResult(BaseModel, Generic[DATA]):
         default=None,
     )
     link: Url | None = Field(description="Link to more information on the outcome of the action", default=None)
+    task_id: str | None = Field(description="The celery task id if the action is pending.", default=None)
 
     @model_validator(mode="after")
     def validate_model(self: Self, info: ValidationInfo) -> Self:  # noqa: C901
@@ -262,8 +270,11 @@ class ActionResult(BaseModel, Generic[DATA]):
         Returns:
             Self: The validated model.
         """
-        if not self.format and self.outcome != "failure":
-            raise ClueValueError("You must set a format if outcome is not failure.")
+        if not self.format and self.outcome == "success":
+            raise ClueValueError("You must set a format if outcome is success.")
+
+        if not self.task_id and self.outcome == "pending":
+            raise ClueValueError("task_id must be set if outcome is pending.")
 
         if self.format == "pivot" and (not self.output or not isinstance(self.output, Url)):
             if isinstance(self.output, str):

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/plugin/__init__.py

@@ -32,11 +32,14 @@ from clue.models.actions import (
     ActionBase,
     ActionResult,
     ActionSpec,
+    ActionStatusRequest,
     ExecuteRequest,
 )
+from clue.models.config import Config
 from clue.models.fetchers import FetcherDefinition, FetcherResult
 from clue.models.network import QueryEntry
 from clue.models.selector import Selector
+from clue.plugin.celery_app import celery_init_app
 from clue.plugin.helpers.token import get_username
 from clue.plugin.models import BulkEntry
 from clue.plugin.utils import Params
@@ -52,6 +55,7 @@ OVERRIDABLE_FUNCTIONS = [
     "liveness",  # Kubernetes liveness probe endpoint
     "readiness",  # Kubernetes readiness probe endpoint
     "run_action",  # Function to execute plugin actions
+    "get_status",  # Function to check the status or result of a pending action
     "run_fetcher",  # Function to execute plugin fetchers
     "setup_actions",  # Runtime action definition generation
     "validate_token",  # Custom authentication token validation
@@ -130,6 +134,37 @@ def build_default_logger() -> logging.Logger:
     return logger
 
 
+config: Config = Config()
+
+
+def create_app(app_name: str, enable_celery: bool = False, tasks: list[str] | None = None):
+    """helper function to create the flask app and set up the celery config if enabled
+
+    Args:
+        enable_celery (bool): whether or not to enable celery
+
+    Returns:
+        _type_: flask app
+    """
+    app = Flask(__name__.split(".")[0])
+    if enable_celery:
+        redis_url = (
+            f"redis://:{config.core.redis.password}@{config.core.redis.host}:{config.core.redis.port}"
+            if config.core.redis.password
+            else f"redis://{config.core.redis.host}:{config.core.redis.port}"
+        )
+        app.config.from_mapping(
+            CELERY=dict(
+                broker_url=redis_url,
+                result_backend=redis_url,
+                result_backend_transport_options={"global_keyprefix": app_name + "_results"},
+                result_expires=3600,  # expire results after one hour
+            ),
+        )
+        celery_init_app(app, tasks)
+    return app
+
+
 class CluePlugin:
     """Helper class for creating clue plugins with proper server responses and behaviour.
 
@@ -259,6 +294,13 @@ class CluePlugin:
     user parameters) that instance will be passed instead, and casting the argument will be necessary.
     """
 
+    get_status: Callable[[Action, ActionStatusRequest, str | None], ActionResult] | None
+    """The function to get the status and result of running actions.
+
+    Accepts the selected action definition as well as an ActionStatusRequest instance which contains the
+    specific task id to check the status of.
+    """
+
     fetchers: list[FetcherDefinition] | None
     "A list of fetcher definitions this plugin supports."
 
@@ -275,6 +317,9 @@ class CluePlugin:
     readiness: Callable[[], Response]
     "A readiness probe for kubernetes implementations of clue."
 
+    enable_celery: bool
+    "Flag to enable celery in plugin"
+
     def __init__(
         self: Self,
         app_name: str,
@@ -284,6 +329,7 @@ class CluePlugin:
         classification: str | None = os.environ.get("CLASSIFICATION", None),
         enable_apm: bool = False,
         enable_cache: Union[bool, Literal["redis"], Literal["local"]] = True,
+        enable_celery: bool = False,
         enrich: Callable[[str, str, Params, str | None], Union[list[QueryEntry], QueryEntry]] | None = None,
         fetchers: list[FetcherDefinition] | None = None,
         liveness: Callable[[], Response] = default_liveness,
@@ -291,6 +337,7 @@ class CluePlugin:
         logger: logging.Logger | None = None,
         readiness: Callable[[], Response] = default_readiness,
         run_action: Callable[[Action, ExecuteRequest, str | None], ActionResult] | None = None,
+        get_status: Callable[[Action, ActionStatusRequest, str | None], ActionResult] | None = None,
         run_fetcher: Callable[[FetcherDefinition, Selector, str | None], FetcherResult] | None = None,
         setup_actions: Callable[[list[Action], str | None], list[Action]] | None = None,
         supported_types: set[str] | str | None = None,
@@ -360,7 +407,7 @@ class CluePlugin:
         """
         self.alternate_bulk_lookup = alternate_bulk_lookup
         # Create Flask app using the module name (before first dot) as app name
-        self.app = Flask(__name__.split(".")[0])
+        self.app = create_app(app_name, enable_celery)
         self.app_name = app_name
 
         # Classification is required for security - must be specified via env var or parameter
@@ -389,7 +436,9 @@ class CluePlugin:
         self.logger = logger if logger else build_default_logger()
 
         self.enrich = enrich
+        self.enable_celery = enable_celery
         self.run_action = run_action
+        self.get_status = get_status
         self.validate_token = validate_token
 
         self.fetchers = fetchers
@@ -662,6 +711,12 @@ class CluePlugin:
         self.app.add_url_rule(
             "/actions/<action_id>/", self.execute_action.__name__, self.execute_action, methods=["POST"]
         )
+        self.app.add_url_rule(
+            "/actions/<action_id>/status/<task_id>",
+            self.get_action_status.__name__,
+            self.get_action_status,
+            methods=["GET"],
+        )
         self.app.add_url_rule("/fetchers/", self.get_fetchers.__name__, self.get_fetchers, methods=["GET"])
         self.app.add_url_rule(
             "/fetchers/<fetcher_id>", self.execute_fetcher.__name__, self.execute_fetcher, methods=["POST"]
@@ -1128,6 +1183,88 @@ class CluePlugin:
 
         return self.make_api_response(result)
 
+    def get_action_status(self: Self, action_id: str, task_id: str):  # noqa: C901
+        """Retrieves the status of the specified action.
+
+        Args:
+            action_id (str): The ID of the action to get the status for
+            task_id (str): The celery task id to get the result from
+
+        Returns:
+            Response: A Response object with an ActionResult as the body.
+        """
+        if not task_id:
+            return self.make_api_response(
+                {}, err="task id not provided. task id is required for this request.", status_code=400
+            )
+
+        if not self.get_status:
+            return self.make_api_response(
+                {}, err=f"{self.app_name} does not support the get action status functions.", status_code=400
+            )
+
+        try:
+            actions = self.__check_actions()
+        except Exception:
+            self.logger.exception("Exception on setup actions:")
+
+            return self.make_api_response({}, err="Error on action setup.", status_code=500)
+
+        if actions is None:
+            actions = self.actions or []
+
+        action_to_check = next((action for action in actions if action.id == action_id), None)
+        if not action_to_check:
+            return self.make_api_response({}, err="Action does not exist", status_code=404)
+
+        token: str | None = None
+        if self.validate_token:
+            self.logger.debug("Executing plugin-provided token validator")
+
+            token, error = self.validate_token()
+
+            if error:
+                return self.make_api_response(None, f"Error on token validation: {error}", status_code=401)
+
+            self.logger.debug("Token is valid")
+        else:
+            self.logger.warning("No token validation provided. The access token will not be provided to the action.")
+
+        try:
+            # Validate request body against the action's parameter schema
+            status_request = ActionStatusRequest(task_id=task_id)
+
+            self.logger.info(
+                "Getting status for Action '%s' with task_id: %s",
+                action_id,
+                task_id,
+            )
+
+            result = self.get_status(action_to_check, status_request, token)
+        except json.JSONDecodeError as e:
+            self.logger.warning("JSON decoding error while getting status: %s", str(e))
+
+            result = ActionResult(
+                outcome="failure",
+                summary=f"Invalid request format. Request body must be valid JSON. Error: {str(e)}",
+            )
+        except ValidationError as err:
+            self.logger.warning("Validation error during execution: %s", str(err))
+
+            result = ActionResult(outcome="failure", summary=f"Validation error: {str(err)}")
+        except ClueException as e:
+            self.logger.exception("ClueException during execution:")
+
+            result = ActionResult(outcome="failure", summary=f"Error encountered during execution: {e.message}")
+        except Exception as e:
+            self.logger.exception("%s during execution:", e.__class__.__name__)
+
+            result = ActionResult(outcome="failure", summary=f"An unknown error occurred during execution: {str(e)}")
+
+        self.logger.info("Action status: %s", result.outcome)
+
+        return self.make_api_response(result)
+
     def get_fetchers(self: Self) -> Response:
         """Get all available fetchers for this plugin.
 

clue_api-1.3.0.dev102/clue/plugin/celery_app.py

@@ -0,0 +1,29 @@
+from typing import List
+
+from celery import Celery, Task
+from flask import Flask
+
+celery = Celery("app")
+
+
+def celery_init_app(app: Flask, tasks: List[str] | None = None) -> None:
+    """initialize the celery worker for the flask app
+
+    Args:
+        app (Flask): flask app instance
+    """
+
+    class FlaskTask(Task):
+        def __call__(self, *args: object, **kwargs: object) -> object:
+            with app.app_context():
+                return self.run(*args, **kwargs)
+
+    celery_app = Celery(app.name, task_cls=FlaskTask)
+    celery_app.config_from_object(app.config["CELERY"])
+    celery_app.set_default()
+    if tasks:
+        celery_app.autodiscover_tasks(
+            tasks,
+            force=True,
+        )
+    app.extensions["celery"] = celery_app

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/clue/services/action_service.py

@@ -184,3 +184,61 @@ def execute_action(plugin_id: str, action_id: str, user: dict[str, Any]) -> Acti
         raise ClueException(
             f"Something went wrong when retrieving the result from plugin '{plugin_id}': {err.__class__.__name__}."
         )
+
+
+def get_action_status(plugin_id: str, action_id: str, task_id: str, user: dict[str, Any]) -> ActionResult:
+    """Gets the status of a specified action with task_id.
+
+    Args:
+        plugin_id (str): The ID of the plugin.
+        action_id (str): The ID of the action to run.
+        task_id (str): The celery task id to fetch the status for
+        user (dict[str, Any]): The user dict of the user running the action.
+
+    Raises:
+        NotFoundException: Raised whenever the plugin or the action doesn't exist.
+        ClueException: Raised whenever an error is returned by the plugin endpoint.
+
+    Returns:
+        ActionResult: The result of the action.
+    """
+    plugin = next((source for source in config.api.external_sources if source.name == plugin_id), None)
+
+    if not plugin:
+        raise NotFoundException(f"Plugin {plugin_id} does not exist.")
+
+    access_token = request.headers.get("Authorization", type=str)
+    if access_token:
+        access_token = access_token.split(" ")[1]
+
+    obo_access_token = None
+    if access_token:
+        obo_access_token, error = auth_service.check_obo(plugin, access_token, user["uname"])
+
+        if error:
+            logger.error("%s: %s", plugin.name, error)
+            return ActionResult(outcome="failure", summary="Invalid token provided.")
+
+    headers = generate_headers(obo_access_token or access_token, access_token if obo_access_token else None)
+
+    try:
+        req_url = urljoin(plugin.url, f"actions/{action_id}/status/{task_id}")
+        logger.debug("Getting status for action %s with task_id %s for user %s", req_url, task_id, user["uname"])
+
+        response = requests.get(
+            req_url,
+            headers=headers,
+            timeout=request.args.get("max_timeout", plugin.default_timeout, type=float),
+        )
+
+        result = response.json()
+
+        if not response.ok:
+            raise ClueException(result["api_error_message"])
+
+        return ActionResult.model_validate(result["api_response"])
+    except (JSONDecodeError, exceptions.ConnectionError) as err:
+        logger.exception(f"Something went wrong when retrieving the status from plugin '{plugin_id}'")
+        raise ClueException(
+            f"Something went wrong when retrieving the status from plugin '{plugin_id}': {err.__class__.__name__}."
+        )

{clue_api-1.3.0.dev92 → clue_api-1.3.0.dev102}/pyproject.toml

@@ -142,7 +142,7 @@ log_cli_level = "WARN"
 [tool.poetry]
 package-mode = true
 name = "clue-api"
-version = "1.3.0.dev92"
+version = "1.3.0.dev102"
 description = "Clue distributed enrichment service"
 authors = ["Canadian Centre for Cyber Security <contact@cyber.gc.ca>"]
 license = "MIT"
@@ -204,6 +204,7 @@ authlib = { version = "<2.0.0", optional = true }
 flask-cors = { version = ">=4.0.1,<7.0.0", optional = true }
 flasgger = { version = "^0.9.7.1", optional = true }
 trino = "^0.336.0"
+celery = "^5.6.2"
 
 [tool.poetry.extras]
 server = [