howler-api 3.4.0.dev912__tar.gz → 3.4.0.dev929__tar.gz

{howler_api-3.4.0.dev912 → howler_api-3.4.0.dev929}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: howler-api
-Version: 3.4.0.dev912
+Version: 3.4.0.dev929
 Summary: Howler - API server
 License: MIT
 Keywords: howler,alerting,gc,canada,cse-cst,cse,cst,cyber,cccs

howler_api-3.4.0.dev929/howler/cronjobs/retention.py

@@ -0,0 +1,129 @@
+import os
+from datetime import datetime, timedelta
+from typing import Any
+
+from apscheduler.schedulers.base import BaseScheduler
+from apscheduler.triggers.cron import CronTrigger
+from pytz import timezone
+
+from howler.common.logging import get_logger
+from howler.config import DEBUG, config
+from howler.datastore.howler_store import HowlerDatastore
+
+logger = get_logger(__file__)
+
+
+def execute():
+    """Delete any hits older than the configured time"""
+    from howler.common.loader import datastore
+
+    delta_kwargs = {str(config.system.retention.limit_unit): config.system.retention.limit_amount}
+
+    cutoff = (datetime.now() - timedelta(**delta_kwargs)).strftime("%Y-%m-%d")
+
+    logger.debug("Removing hits older than %s", cutoff)
+
+    ds = datastore()
+
+    ds.hit.delete_by_query(f"event.created:{{* TO {cutoff}}} OR howler.expiry:{{* TO now}}")
+
+    ds.hit.commit()
+
+    logger.debug("Deletion complete")
+
+    logger.debug("Cleaning analytics with no matching hits")
+    _remove_analytics_without_hits(ds)
+
+
+def _remove_analytics_without_hits(ds: HowlerDatastore):
+
+    matched_analytics = _find_analytics_with_hits(ds)
+
+    if matched_analytics is not None:
+        ds.analytic.delete_by_search_object(
+            {
+                "bool": {
+                    "filter": [
+                        {
+                            "bool": {
+                                "must_not": [
+                                    {"exists": {"field": "rule"}},
+                                    {"exists": {"field": "rule_type"}},
+                                ]
+                            }
+                        }
+                    ],
+                    "must_not": [{"terms": {"name": matched_analytics}}],
+                }
+            }
+        )
+        ds.analytic.commit()
+    else:
+        logger.warning(
+            "Aggregation search for matched analytics did not run or returned no results. "
+            "There is likely an issue with the query. Skipping cleanup."
+        )
+
+
+def _find_analytics_with_hits(ds: HowlerDatastore) -> list[str] | None:
+
+    total_analytics = ds.analytic.count("id:*", filters=None)["count"]
+
+    if total_analytics:
+        matched_analytics = ds.hit.search(
+            "howler.id:*",
+            aggregations=[
+                (
+                    "matched_analytics",
+                    {
+                        "terms": {
+                            "field": "howler.analytic",
+                            "size": total_analytics,
+                        }
+                    },
+                )
+            ],
+            rows=0,
+        )
+
+        if "matched_analytics" in matched_analytics["aggregations"]:
+            matched_analytic_names = [
+                bucket["key"] for bucket in matched_analytics["aggregations"]["matched_analytics"]["buckets"]
+            ]
+        else:
+            return None
+
+    else:
+        matched_analytic_names = []
+
+    return matched_analytic_names
+
+
+def setup_job(sched: BaseScheduler):
+    """Initialize the retention job"""
+    if not config.system.retention.enabled:
+        if not DEBUG or config.system.type == "production":
+            logger.warning("Retention cronjob disabled! This is not recommended for a production settings.")
+
+        return
+
+    logger.debug("Initializing retention cronjob with cron %s", config.system.retention.crontab)
+
+    if DEBUG:
+        _kwargs: dict[str, Any] = {"next_run_time": datetime.now()}
+    else:
+        _kwargs = {}
+
+    if sched.get_job("retention"):
+        logger.debug("Retention job already running!")
+        return
+
+    sched.add_job(
+        id="retention",
+        func=execute,
+        trigger=CronTrigger.from_crontab(
+            config.system.retention.crontab, timezone=timezone(os.getenv("SCHEDULER_TZ", "America/Toronto"))
+        ),
+        **_kwargs,
+    )
+    logger.debug("Initialization complete")

{howler_api-3.4.0.dev912 → howler_api-3.4.0.dev929}/howler/datastore/collection.py

@@ -38,7 +38,7 @@ from howler.datastore.support.schemas import (
     default_index,
     default_mapping,
 )
-from howler.datastore.types import SearchResult
+from howler.datastore.types import AggSearchResult, SearchResult
 from howler.odm.base import (
     BANNED_FIELDS,
     IP,
@@ -211,9 +211,11 @@ class ESCollection(Generic[ModelType]):
         "sort": DEFAULT_SORT,
         "df": None,
         "script_fields": [],
+        "aggregations": None,
     }
     IGNORE_ENSURE_COLLECTION: bool = False
     ENSURE_COLLECTION_WARNED: bool = False
+    CUSTOM_AGG_PREFIX: str = "_custom_agg__"
 
     def __init__(self, datastore: ESStore, name, model_class=None, validate=True, max_attempts=10):
         self.replicas = int(
@@ -1119,7 +1121,7 @@ class ESCollection(Generic[ModelType]):
         except elasticsearch.NotFoundError:
             return False
 
-    def delete_by_query(self, query, workers=20, sort=None, max_docs=None):
+    def delete_by_query(self, query: str, sort=None, max_docs=None):
         """This function should delete the underlying documents referenced by the query.
         It should return true if the documents were in fact properly deleted.
 
@@ -1127,7 +1129,18 @@ class ESCollection(Generic[ModelType]):
         :param workers: Number of workers used for deletion if basic currency delete is used
         :return: True is delete successful
         """
-        query = {"bool": {"must": {"query_string": {"query": query}}}}
+        query_obj = {"bool": {"must": {"query_string": {"query": query}}}}
+        success = self.delete_by_search_object(query=query_obj, sort=sort, max_docs=max_docs)
+        return success
+
+    def delete_by_search_object(self, query: dict, sort=None, max_docs=None):
+        """Delete the underlying documents matching the query object.
+        Returns true if the documents were in fact properly deleted.
+
+        :param query: Query object following elasticsearch request structure
+        :param workers: Number of workers used for deletion if basic currency delete is used
+        :return: True is delete successful
+        """
         info = self._delete_async(self.name, query=query, sort=sort_str(parse_sort(sort)), max_docs=max_docs)
         return info.get("deleted", 0) != 0
 
@@ -1503,6 +1516,26 @@ class ESCollection(Generic[ModelType]):
                 },
             }
 
+        # Add any arbitrary aggregations
+        if parsed_values["aggregations"]:
+            query_body.setdefault("aggregations", {})
+            cluster_settings = self.datastore.client.cluster.get_settings(include_defaults=True, flat_settings=True)
+            flattened_settings = {
+                **cluster_settings["defaults"],
+                **cluster_settings["transient"],
+                **cluster_settings["persistent"],
+            }
+            max_buckets = int(flattened_settings["search.max_buckets"])
+            for agg_name, agg_args in parsed_values["aggregations"]:
+                if any("size" in agg_def and agg_def["size"] > max_buckets for agg_def in agg_args.values()):
+                    # verify the size of the agg query doesn't exceed the max
+                    warnings.warn(
+                        f"Aggregation {agg_name} has a size argument higher than the maximum allowed "
+                        f"buckets of the cluster ({max_buckets}). Skipping aggregation."
+                    )
+                    continue
+                query_body["aggregations"][f"{self.CUSTOM_AGG_PREFIX}{agg_name}"] = agg_args
+
         try:
             if deep_paging_id is not None and not deep_paging_id == "*":
                 # Get the next page
@@ -1554,6 +1587,8 @@ class ESCollection(Generic[ModelType]):
         use_archive: bool = False,
         track_total_hits: bool = False,
         script_fields: list[str] = [],
+        *,
+        aggregations: None = None,
     ) -> SearchResult[ModelType]: ...
 
     @overload
@@ -1572,8 +1607,50 @@ class ESCollection(Generic[ModelType]):
         use_archive: bool = False,
         track_total_hits: bool = False,
         script_fields: list[str] = [],
+        *,
+        aggregations: None = None,
     ) -> SearchResult[dict[str, typing.Any]]: ...
 
+    @overload
+    def search(
+        self,
+        query: str | None,
+        as_obj: Literal[True] = True,
+        offset: int = 0,
+        rows: int | None = None,
+        sort: typing.Any = None,
+        fl: str | None = None,
+        timeout: int | None = None,
+        filters: list[str] | str | None = None,
+        access_control: typing.Any = None,
+        deep_paging_id: str | None = None,
+        use_archive: bool = False,
+        track_total_hits: bool = False,
+        script_fields: list[str] = [],
+        *,
+        aggregations: list[tuple[str, dict]],
+    ) -> AggSearchResult[ModelType]: ...
+
+    @overload
+    def search(
+        self,
+        query: str | None,
+        as_obj: Literal[False],
+        offset: int = 0,
+        rows: int | None = None,
+        sort: typing.Any = None,
+        fl: str | None = None,
+        timeout: int | None = None,
+        filters: list[str] | str | None = None,
+        access_control: typing.Any = None,
+        deep_paging_id: str | None = None,
+        use_archive: bool = False,
+        track_total_hits: bool = False,
+        script_fields: list[str] = [],
+        *,
+        aggregations: list[tuple[str, dict]],
+    ) -> AggSearchResult[dict[str, typing.Any]]: ...
+
     def search(
         self,
         query,
@@ -1589,6 +1666,8 @@ class ESCollection(Generic[ModelType]):
         use_archive=False,
         track_total_hits=None,
         script_fields=[],
+        *,
+        aggregations=None,
     ):
         """This function should perform a search through the datastore and return a
         search result object that consist on the following::
@@ -1605,6 +1684,14 @@ class ESCollection(Generic[ModelType]):
                     }, ...]
             }
 
+        If aggregations are provided the search result will include an additional field::
+
+            {
+                "aggregations": {       # Dictionary where the keys are the keys of the `aggregations` parameter
+                    "agg_name": {...}   #   and the values are the results of the aggregations
+                }
+            }
+
         :param script_fields: List of name/script tuple of fields to be evaluated at runtime
         :param track_total_hits: Return to total matching document count
         :param use_archive: Query also the archive
@@ -1618,6 +1705,8 @@ class ESCollection(Generic[ModelType]):
         :param timeout: maximum time of execution
         :param filters: additional queries to run on the original query to reduce the scope
         :param access_control: access control parameters to limiti the scope of the query
+        :param aggregations: optional list of arbitrary aggregations to run alongside the query
+            structured the same way as the es rest query aggs field
         :return: a search result object
         """
         if offset is None:
@@ -1660,6 +1749,9 @@ class ESCollection(Generic[ModelType]):
         if script_fields:
             args.append(("script_fields", script_fields))
 
+        if aggregations:
+            args.append(("aggregations", aggregations))
+
         result = self._search(
             args,
             deep_paging_id=deep_paging_id,
@@ -1667,13 +1759,24 @@ class ESCollection(Generic[ModelType]):
             track_total_hits=track_total_hits,
         )
 
-        ret_data: SearchResult = {
+        ret_data: SearchResult | AggSearchResult
+        ret_data = {
             "offset": int(offset),
             "rows": int(rows),
             "total": int(result["hits"]["total"]["value"]),
             "items": [self._format_output(doc, field_list, as_obj=as_obj) for doc in result["hits"]["hits"]],
         }
 
+        if aggregations:
+            ret_data = {
+                **ret_data,
+                "aggregations": {
+                    k[len(self.CUSTOM_AGG_PREFIX) :]: v
+                    for k, v in result.get("aggregations", {}).items()
+                    if k.startswith(self.CUSTOM_AGG_PREFIX)
+                },
+            }
+
         new_deep_paging_id = result.get("_scroll_id", None)
 
         # Check if the scroll is finished and close it

{howler_api-3.4.0.dev912 → howler_api-3.4.0.dev929}/howler/datastore/types.py

@@ -16,3 +16,7 @@ class SearchResult(TypedDict, Generic[SearchResultType]):
     total: int
     items: list[SearchResultType]
     next_deep_paging_id: NotRequired[str | None]
+
+
+class AggSearchResult(SearchResult[SearchResultType]):
+    aggregations: dict[str, dict]

{howler_api-3.4.0.dev912 → howler_api-3.4.0.dev929}/pyproject.toml

@@ -152,7 +152,7 @@ suppress-none-returning = true
 [tool.poetry]
 package-mode = true
 name = "howler-api"
-version = "3.4.0.dev912"
+version = "3.4.0.dev929"
 description = "Howler - API server"
 authors = [
     "Canadian Centre for Cyber Security <howler@cyber.gc.ca>",

howler_api-3.4.0.dev912/howler/cronjobs/retention.py

@@ -1,61 +0,0 @@
-import os
-from datetime import datetime, timedelta
-from typing import Any
-
-from apscheduler.schedulers.base import BaseScheduler
-from apscheduler.triggers.cron import CronTrigger
-from pytz import timezone
-
-from howler.common.logging import get_logger
-from howler.config import DEBUG, config
-
-logger = get_logger(__file__)
-
-
-def execute():
-    """Delete any hits older than the configured time"""
-    from howler.common.loader import datastore
-
-    delta_kwargs = {str(config.system.retention.limit_unit): config.system.retention.limit_amount}
-
-    cutoff = (datetime.now() - timedelta(**delta_kwargs)).strftime("%Y-%m-%d")
-
-    logger.debug("Removing hits older than %s", cutoff)
-
-    ds = datastore()
-
-    ds.hit.delete_by_query(f"event.created:{{* TO {cutoff}}} OR howler.expiry:{{* TO now}}")
-
-    ds.hit.commit()
-
-    logger.debug("Deletion complete")
-
-
-def setup_job(sched: BaseScheduler):
-    """Initialize the retention job"""
-    if not config.system.retention.enabled:
-        if not DEBUG or config.system.type == "production":
-            logger.warning("Retention cronjob disabled! This is not recommended for a production settings.")
-
-        return
-
-    logger.debug("Initializing retention cronjob with cron %s", config.system.retention.crontab)
-
-    if DEBUG:
-        _kwargs: dict[str, Any] = {"next_run_time": datetime.now()}
-    else:
-        _kwargs = {}
-
-    if sched.get_job("retention"):
-        logger.debug("Retention job already running!")
-        return
-
-    sched.add_job(
-        id="retention",
-        func=execute,
-        trigger=CronTrigger.from_crontab(
-            config.system.retention.crontab, timezone=timezone(os.getenv("SCHEDULER_TZ", "America/Toronto"))
-        ),
-        **_kwargs,
-    )
-    logger.debug("Initialization complete")