# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one # or more contributor license agreements. Licensed under the Elastic License # 2.0; you may not use this file except in compliance with the Elastic License # 2.0. import datetime from typing import List, Optional, Type from .connector import Kibana DEFAULT_PAGE_SIZE = 10 class BaseResource(dict): BASE_URI = "" ID_FIELD = "id" @property def id(self): return self.get(self.ID_FIELD) @classmethod def bulk_create(cls, resources: list): for r in resources: assert isinstance(r, cls) responses = Kibana.current().post(cls.BASE_URI + "/_bulk_create", data=resources) return [cls(r) for r in responses] def create(self): response = Kibana.current().post(self.BASE_URI, data=self) self.update(response) return self @classmethod def find(cls, per_page=None, **params) -> iter: if per_page is None: per_page = DEFAULT_PAGE_SIZE params.setdefault("sort_field", "_id") params.setdefault("sort_order", "asc") return ResourceIterator(cls, cls.BASE_URI + "/_find", per_page=per_page, **params) @classmethod def from_id(cls, resource_id) -> 'BaseResource': return Kibana.current().get(cls.BASE_URI, params={cls.ID_FIELD: resource_id}) def put(self): response = Kibana.current().put(self.BASE_URI, data=self.to_dict()) self._update_from(response) return self def delete(self): return Kibana.current().delete(self.BASE_URI, params={"id": self.id}) class ResourceIterator(object): def __init__(self, cls: Type[BaseResource], uri: str, per_page: int, **params: dict): self.cls = cls self.uri = uri self.params = params self.page = 0 self.per_page = per_page self.fetched = 0 self.current = None self.total = None self.batch = [] self.batch_pos = 0 self.kibana = Kibana.current() def __iter__(self): return self def _batch(self): params = dict(per_page=self.per_page, page=self.page + 1, **self.params) response = self.kibana.get(self.uri, params=params, error=True) self.page = response["page"] self.per_page = response["perPage"] self.total = response["total"] self.batch = response["data"] self.batch_pos = 0 self.fetched += len(self.batch) def __next__(self) -> BaseResource: if self.total is None or 0 < self.batch_pos == len(self.batch) == self.per_page: self._batch() if self.batch_pos < len(self.batch): result = self.cls(self.batch[self.batch_pos]) self.batch_pos += 1 return result raise StopIteration() class RuleResource(BaseResource): BASE_URI = "/api/detection_engine/rules" @staticmethod def _add_internal_filter(is_internal: bool, params: dict) -> dict: custom_filter = f'alert.attributes.tags:"__internal_immutable:{str(is_internal).lower()}"' if params.get("filter"): params["filter"] = f"({params['filter']}) and ({custom_filter})" else: params["filter"] = custom_filter return params @classmethod def find_custom(cls, **params): params = cls._add_internal_filter(False, params) return cls.find(**params) @classmethod def find_elastic(cls, **params): # GET params: # * `sort_field` # * `sort_order` # * `filter` (accepts KQL) # alert.attributes.name:mshta # alert.attributes.enabled:true/false # # ... # i.e. Rule.find_elastic(filter="alert.attributes.name:mshta") params = cls._add_internal_filter(True, params) return cls.find(**params) def put(self): # id and rule_id are mutually exclusive rule_id = self.get("rule_id") self.pop("rule_id", None) try: # apparently Kibana doesn't like `rule_id` for existing documents return super(RuleResource, self).update() except Exception: # if it fails, restore the id back if rule_id: self["rule_id"] = rule_id raise class Signal(BaseResource): BASE_URI = "/api/detection_engine/signals" def __init__(self): raise NotImplementedError("Signals can't be instantiated yet") @classmethod def search(cls, query_dsl: dict, size: Optional[int] = 10): payload = dict(size=size, **query_dsl) return Kibana.current().post(f"{cls.BASE_URI}/search", data=payload) @classmethod def last_signal(cls) -> (int, datetime.datetime): query_dsl = { "aggs": { "lastSeen": {"max": {"field": "@timestamp"}} }, 'query': { "bool": { "filter": [ {"match": {"signal.status": "open"}} ] } }, "size": 0, "track_total_hits": True } response = cls.search(query_dsl) last_seen = response.get("aggregations", {}).get("last_seen", {}).get("value_as_string") num_signals = response.get("hits", {}).get("total", {}).get("value") if last_seen is not None: last_seen = datetime.datetime.strptime(last_seen, "%Y-%m-%dT%H:%M:%S.%f%z") return num_signals, last_seen @classmethod def all(cls, size: Optional[int] = 10): return cls.search({"query": {"bool": {"filter": {"match_all": {}}}}}, size=size) @classmethod def set_status_many(cls, signal_ids: List[str], status: str) -> dict: return Kibana.current().post(f"{cls.BASE_URI}/status", data={"signal_ids": signal_ids, "status": status}) @classmethod def close_many(cls, signal_ids: List[str]): return cls.set_status_many(signal_ids, "closed") @classmethod def open_many(cls, signal_ids: List[str]): return cls.set_status_many(signal_ids, "open")