Refresh Kibana module with API updates

Author: brokensound77

detection_rules/kbwrap.py

@@ -53,7 +53,7 @@ def upload_rule(ctx, rules, replace_id):
         api_payloads.append(rule)
 
     with kibana:
-        results = RuleResource.bulk_create(api_payloads)
+        results = RuleResource.bulk_create_legacy(api_payloads)
 
     success = []
     errors = []

kibana/connector.py

@@ -10,6 +10,7 @@
 import sys
 import threading
 import uuid
+from typing import List, Optional, Union
 
 import requests
 from elasticsearch import Elasticsearch
@@ -72,6 +73,16 @@ def version(self):
         if self.status:
             return self.status.get("version", {}).get("number")
 
+    @staticmethod
+    def ndjson_file_data_prep(lines: List[dict], filename: str) -> (dict, str):
+        """Prepare a request for an ndjson file upload to Kibana."""
+        data = ('\n'.join(json.dumps(r) for r in lines) + '\n')
+        boundary = '----JustAnotherBoundary'
+        bounded_data = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{filename}"\r\n'
+                        f'Content-Type: application/x-ndjson\r\n\r\n{data}\r\n--{boundary}--\r\n').encode('utf-8')
+        headers = {'content-type': f'multipart/form-data; boundary={boundary}'}
+        return headers, bounded_data
+
     def url(self, uri):
         """Get the full URL given a URI."""
         assert self.kibana_url is not None
@@ -81,15 +92,20 @@ def url(self, uri):
             uri = "s/{}/{}".format(self.space, uri)
         return f"{self.kibana_url}/{uri}"
 
-    def request(self, method, uri, params=None, data=None, error=True, verbose=True, raw=False, **kwargs):
+    def request(self, method, uri, params=None, data=None, raw_data=None, error=True, verbose=True, raw=False,
+                **kwargs) -> Optional[Union[requests.Response, dict]]:
         """Perform a RESTful HTTP request with JSON responses."""
-        params = params or {}
         url = self.url(uri)
-        params = {k: v for k, v in params.items()}
+        params = params or {}
+        params = json.dumps(params)
         body = None
         if data is not None:
             body = json.dumps(data)
 
+        assert not (body and raw_data), "Cannot provide both data and raw_data"
+
+        body = body or raw_data
+
         response = self.session.request(method, url, params=params, data=body, **kwargs)
 
         if response.status_code != 200:
@@ -100,14 +116,16 @@ def request(self, method, uri, params=None, data=None, error=True, verbose=True,
             try:
                 response.raise_for_status()
             except requests.exceptions.HTTPError:
+                if response.status_code == 404:
+                    raise NotImplementedError(f'API endpoint {uri} not implemented for Kibana version {self.version}')
                 if verbose:
                     print(response.content.decode("utf-8"), file=sys.stderr)
                 raise
 
         if not response.content:
             return
 
-        return response.content if raw else response.json()
+        return response if raw else response.json()
 
     def get(self, uri, params=None, data=None, error=True, **kwargs):
         """Perform an HTTP GET."""

kibana/defenitions.py

@@ -0,0 +1,44 @@
+# 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.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Literal, NewType, Union
+
+
+RuleBulkActions = Literal['enable', 'disable', 'delete', 'duplicate', 'export', 'edit']
+
+RuleBulkAddTags = NewType('RuleBulkAddTags', List[str])
+RuleBulkDeleteTags = NewType('RuleBulkDeleteTags', List[str])
+RuleBulkSetTags = NewType('RuleBulkSetTags', List[str])
+RuleBulkAddIndexPatterns = NewType('RuleBulkAddIndexPatterns', List[str])
+RuleBulkDeleteIndexPatterns = NewType('RuleBulkDeleteIndexPatterns', List[str])
+RuleBulkSetIndexPatterns = NewType('RuleBulkSetIndexPatterns', List[str])
+RuleBulkSetTimelineTitle = NewType('RuleBulkSetTimelineTitle', Dict[Literal['timeline_id', 'timeline_title'], str])
+RuleBulkSetSchedule = NewType('RuleBulkSetSchedule', Dict[Literal['interval', 'lookback'], str])
+RuleBulkAddRuleActions = NewType('RuleBulkAddRuleActions', Dict[Literal['actions', 'throttle'], Union[List[dict], dict]])
+RuleBulkSetRuleActions = NewType('RuleBulkSetRuleActions', Dict[Literal['actions', 'throttle'], Union[List[dict], dict]])
+
+RuleBulkEditActionTypes = Union[
+    RuleBulkAddTags,
+    RuleBulkDeleteTags,
+    RuleBulkSetTags,
+    RuleBulkAddIndexPatterns,
+    RuleBulkDeleteIndexPatterns,
+    RuleBulkSetIndexPatterns,
+    RuleBulkSetTimelineTitle,
+    RuleBulkSetSchedule,
+    RuleBulkAddRuleActions,
+    RuleBulkSetRuleActions
+]
+
+@dataclass
+class RuleBulkEditAction:
+    type: RuleBulkEditActionTypes
+    value: Any
+
+
+@dataclass
+class RuleBulkDuplicateAction:
+    include_exceptions: bool

kibana/resources.py

@@ -6,7 +6,10 @@
 import datetime
 from typing import List, Optional, Type
 
+import json
+
 from .connector import Kibana
+from . import defenitions
 
 DEFAULT_PAGE_SIZE = 10
 
@@ -20,10 +23,12 @@ def id(self):
         return self.get(self.ID_FIELD)
 
     @classmethod
-    def bulk_create(cls, resources: list):
+    def bulk_create_legacy(cls, resources: list):
         for r in resources:
             assert isinstance(r, cls)
 
+        # _bulk_create is being deprecated. Leave for backwards compat only
+        # the new API would be import with multiple rules within an ndjson request
         responses = Kibana.current().post(cls.BASE_URI + "/_bulk_create", data=resources)
         return [cls(r) for r in responses]
 
@@ -127,6 +132,64 @@ def find_elastic(cls, **params):
         params = cls._add_internal_filter(True, params)
         return cls.find(**params)
 
+    @classmethod
+    def bulk_action(cls, action: defenitions.RuleBulkActions, rule_ids: Optional[List[str]] = None,
+                    query: Optional[str] = None, dry_run: Optional[bool] = False,
+                    edit: Optional[defenitions.RuleBulkEditAction] = None,
+                    duplicate: Optional[defenitions.RuleBulkDuplicateAction] = None) -> (dict, List['RuleResource']):
+        assert not (rule_ids and query), 'Cannot provide both rule_ids and query'
+
+        if action == 'edit':
+            assert edit, 'edit action requires edit object'
+
+        params = dict(dry_run=dry_run)
+        data = dict(query=query, ids=rule_ids, action=action, edit=edit, duplicate=duplicate)
+        response = Kibana.current().post(cls.BASE_URI + "/_bulk_action", params=params, data=data)
+
+        results = response['attributes']['results']
+        result_ids = [r['rule_id'] for r in results['updated']]
+        result_ids.extend([r['rule_id'] for r in results['created']])
+        rule_resources = cls.export_rules(result_ids)
+        return response, rule_resources
+
+    @classmethod
+    def bulk_enable(cls, rule_ids: Optional[List[str]] = None,
+                    query: Optional[str] = None, dry_run: Optional[bool] = False) -> (dict, List['RuleResource']):
+        """Bulk enable rules using _bulk_action."""
+        return cls.bulk_action("enable", rule_ids=rule_ids, query=query, dry_run=dry_run)
+
+    @classmethod
+    def bulk_disable(cls, rule_ids: Optional[List[str]] = None,
+                    query: Optional[str] = None, dry_run: Optional[bool] = False) -> (dict, List['RuleResource']):
+        """Bulk disable rules using _bulk_action."""
+        return cls.bulk_action("disable", rule_ids=rule_ids, query=query, dry_run=dry_run)
+
+    @classmethod
+    def bulk_delete(cls,rule_ids: Optional[List[str]] = None,
+                    query: Optional[str] = None, dry_run: Optional[bool] = False) -> (dict, List['RuleResource']):
+        """Bulk delete rules using _bulk_action."""
+        return cls.bulk_action("delete", rule_ids=rule_ids, query=query, dry_run=dry_run)
+
+    @classmethod
+    def bulk_duplicate(cls, rule_ids: Optional[List[str]] = None,
+                    query: Optional[str] = None, dry_run: Optional[bool] = False,
+                    duplicate: Optional[defenitions.RuleBulkDuplicateAction] = None) -> (dict, List['RuleResource']):
+        """Bulk duplicate rules using _bulk_action."""
+        return cls.bulk_action("duplicate", rule_ids=rule_ids, query=query, dry_run=dry_run, duplicate=duplicate)
+
+    @classmethod
+    def bulk_export(cls, rule_ids: Optional[List[str]] = None,
+                    query: Optional[str] = None, dry_run: Optional[bool] = False) -> (dict, List['RuleResource']):
+        """Bulk export rules using _bulk_action."""
+        return cls.bulk_action("export", rule_ids=rule_ids, query=query, dry_run=dry_run)
+
+    @classmethod
+    def bulk_edit(cls, rule_ids: Optional[List[str]] = None,
+                    query: Optional[str] = None, dry_run: Optional[bool] = False,
+                    edit: Optional[defenitions.RuleBulkEditAction] = None) -> (dict, List['RuleResource']):
+        """Bulk edit rules using _bulk_action."""
+        return cls.bulk_action("edit", rule_ids=rule_ids, query=query, dry_run=dry_run, edit=edit)
+
     def put(self):
         # id and rule_id are mutually exclusive
         rule_id = self.get("rule_id")
@@ -142,6 +205,41 @@ def put(self):
 
             raise
 
+    @classmethod
+    def import_rules(cls, rules: List[dict], overwrite: bool = False, overwrite_exceptions: bool = False,
+                     overwrite_action_connectors: bool = False) -> (dict, List['RuleResource']):
+        """Import a list of rules into Kibana using the _import API and return the response and successful imports."""
+        url = f'{cls.BASE_URI}/_import'
+        params = dict(
+            overwrite=overwrite,
+            overwrite_exceptions=overwrite_exceptions,
+            overwrite_action_connectors=overwrite_action_connectors
+        )
+        rule_ids = [r['rule_id'] for r in rules]
+        headers, raw_data = Kibana.ndjson_file_data_prep(rules, "import.ndjson")
+        response = Kibana.current().post(url, headers=headers, params=params, raw_data=raw_data)
+        errors = response.get("errors", [])
+        error_rule_ids = [e['rule_id'] for e in errors]
+
+        # successful rule_ids are not returned, so they must be implicitly inferred from errored rule_ids
+        successful_rule_ids = [r for r in rule_ids if r not in error_rule_ids]
+        rule_resources = cls.export_rules(successful_rule_ids)
+        return response, rule_resources
+
+    @classmethod
+    def export_rules(cls, rule_ids: Optional[List[str]] = None,
+                     exclude_export_details: bool = False) -> List['RuleResource']:
+        """Export a list of rules from Kibana using the _export API."""
+        url = f'{cls.BASE_URI}/_export'
+
+        if rule_ids:
+            rule_ids = {'objects': [{'rule_id': r} for r in rule_ids]}
+
+        params = dict(exclude_export_details=exclude_export_details)
+        response = Kibana.current().post(url, params=params, data=rule_ids, raw=True)
+        data = [json.loads(r) for r in response.text.splitlines()]
+        return [cls(r) for r in data]
+
 
 class Signal(BaseResource):
     BASE_URI = "/api/detection_engine/signals"