Merge pull request #387 from opencybersecurityalliance/develop

v1.7.3

CHANGELOG.rst

@@ -9,6 +9,18 @@ The format is based on `Keep a Changelog`_.
 Unreleased
 ==========
 
+1.7.3 (2023-07-26)
+==================
+
+Added
+-----
+
+- stix-shifter data source interface diagnosis module
+- ``stix-shifter-diag``: stix-shifter data source interface diagnosis utility
+- Docs on ``stix-shifter-diag``
+- Error message update to point to ``stix-shifter-diag``
+- Unit tests of the diagnosis module and CLI utility
+
 1.7.2 (2023-07-26)
 ==================
 

bin/stix-shifter-diag

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+
+import argparse
+from kestrel_datasource_stixshifter.diagnosis import Diagnosis
+
+def gen_patterns():
+    time_range = "START t'2000-01-01T00:00:00.000Z' STOP t'3000-01-01T00:00:00.000Z'"
+    patterns = [
+        "[ipv4-addr:value LIKE '%']",
+        "[process:pid > 0]",
+        "[email-addr:value LIKE '%']",
+    ]
+    return [" ".join([p, time_range]) for p in patterns]
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Kestrel stix-shifter data source interface diagnosis")
+    parser.add_argument("datasource", help="data source name specified in stixshifter.yaml")
+    args = parser.parse_args()
+
+    patterns = gen_patterns()
+
+    diag = Diagnosis(args.datasource)
+
+    diag.diagnose_config()
+    diag.diagnose_ping()
+    diag.diagnose_translate_query(patterns[0])
+    diag.diagnose_run_query_and_retrieval_result(patterns, 1)
+    diag.diagnose_run_query_and_retrieval_result(patterns, 5)

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = kestrel-lang
-version = 1.7.2
+version = 1.7.3
 description = Kestrel Threat Hunting Language
 long_description = file:README.rst
 long_description_content_type = text/x-rst
@@ -24,7 +24,9 @@ project_urls =
 packages = find:
 package_dir =
     =src
-scripts = bin/kestrel
+scripts =
+    bin/kestrel
+    bin/stix-shifter-diag
 python_requires = >= 3.8
 install_requires =
     typeguard>=4.0.0

src/kestrel/exceptions.py

@@ -198,7 +198,7 @@ def __init__(self, uri, itf, msg=""):
 class DataSourceError(KestrelException):
     def __init__(self, error, suggestion=""):
         if not suggestion:
-            suggestion = "please check data source config or test the query manually"
+            suggestion = "please check data source config or diagnose with stix-shifter-diag command"
         super().__init__(
             error,
             suggestion,

src/kestrel/utils.py

@@ -16,6 +16,16 @@ def lowered_str_list(xs):
     return [x.lower() for x in xs if isinstance(x, str)]
 
 
+def mask_value_in_nested_dict(d):
+    if d:
+        for k, v in d.items():
+            if isinstance(v, collections.abc.Mapping):
+                d[k] = mask_value_in_nested_dict(v)
+            elif isinstance(v, str):
+                d[k] = "********"
+    return d
+
+
 def update_nested_dict(dict_old, dict_new):
     if dict_new:
         for k, v in dict_new.items():

src/kestrel_datasource_stixshifter/diagnosis.py

@@ -0,0 +1,149 @@
+import json
+from copy import deepcopy
+from multiprocessing import Queue
+from kestrel.utils import mask_value_in_nested_dict
+from kestrel_datasource_stixshifter.config import (
+    set_stixshifter_logging_level,
+    get_datasource_from_profiles,
+    load_options,
+    load_profiles,
+)
+from kestrel_datasource_stixshifter.worker import STOP_SIGN
+from kestrel_datasource_stixshifter.query import translate_query
+from kestrel_datasource_stixshifter.worker.transmitter import Transmitter
+from stix_shifter.stix_transmission import stix_transmission
+
+
+class Diagnosis:
+    def __init__(self, datasource_name):
+        self.datasource_name = datasource_name
+        self.profiles = load_profiles()
+        self.kestrel_options = load_options()
+        (
+            self.connector_name,
+            self.connection_dict,
+            self.configuration_dict,
+            self.retrieval_batch_size,
+            self.cool_down_after_transmission,
+        ) = get_datasource_from_profiles(datasource_name, self.profiles)
+        self.if_fast_translation = (
+            self.connector_name in self.kestrel_options["fast_translate"]
+        )
+        set_stixshifter_logging_level()
+
+    def diagnose_config(self):
+        print()
+        print()
+        print()
+        print("## Diagnose: config verification")
+
+        configuration_dict_masked = mask_value_in_nested_dict(
+            deepcopy(self.configuration_dict)
+        )
+
+        print()
+        print("#### Kestrel specific config")
+        print(f"retrieval batch size: {self.retrieval_batch_size}")
+        print(f"cool down after transmission: {self.cool_down_after_transmission}")
+        print(f"enable fast translation: {self.if_fast_translation}")
+
+        print()
+        print("#### Config to be passed to stix-shifter")
+        print(f"connector name: {self.connector_name}")
+        print(
+            "connection object [ref: https://github.com/opencybersecurityalliance/stix-shifter/blob/develop/OVERVIEW.md#connection]:"
+        )
+        print(json.dumps(self.connection_dict, indent=4))
+        print(
+            "configuration object [ref: https://github.com/opencybersecurityalliance/stix-shifter/blob/develop/OVERVIEW.md#configuration]:"
+        )
+        print(json.dumps(configuration_dict_masked, indent=4))
+
+    def diagnose_ping(self):
+        print()
+        print()
+        print()
+        print("## Diagnose: stix-shifter to data source connection (network, auth)")
+
+        transmission = stix_transmission.StixTransmission(
+            self.connector_name,
+            self.connection_dict,
+            self.configuration_dict,
+        )
+
+        result = transmission.ping()
+
+        print()
+        print("#### Results from stixshifter transmission.ping()")
+        print(json.dumps(result, indent=4))
+
+    def diagnose_translate_query(self, stix_pattern, quiet=False):
+        if not quiet:
+            print()
+            print()
+            print()
+            print("## Diagnose: stix-shifter query translation")
+
+        dsl = translate_query(
+            self.connector_name,
+            {},
+            stix_pattern,
+            self.connection_dict,
+        )
+
+        if not quiet:
+            print()
+            print("#### Input pattern")
+            print(stix_pattern)
+            print()
+            print("#### Output data source native query")
+            print(json.dumps(dsl, indent=4))
+
+        return dsl
+
+    def diagnose_run_query_and_retrieval_result(self, stix_patterns, max_batch_cnt):
+        print()
+        print()
+        print()
+        print(f"## Diagnose: stix-shifter query execution: <={max_batch_cnt} batch(s)")
+
+        result_queue = Queue()
+        result_counts = []
+
+        for pattern in stix_patterns:
+            for query in self.diagnose_translate_query(pattern, True)["queries"]:
+                transmitter = Transmitter(
+                    self.connector_name,
+                    self.connection_dict,
+                    self.configuration_dict,
+                    self.retrieval_batch_size,
+                    self.cool_down_after_transmission,
+                    query,
+                    result_queue,
+                    max_batch_cnt * self.retrieval_batch_size,
+                )
+
+                transmitter.run()
+                result_queue.put(STOP_SIGN)
+
+                print()
+                print("#### data retrieval results:")
+                for packet in iter(result_queue.get, STOP_SIGN):
+                    if packet.success:
+                        cnt = len(packet.data)
+                        result_counts.append(cnt)
+                        print(f"one batch retrieved: {cnt} entries")
+                    else:
+                        print(packet.log)
+
+                if result_counts:
+                    break
+                else:
+                    print(f"no result matched for query: {query}, go next query")
+
+            if result_counts:
+                break
+            else:
+                print(f"no result matched for pattern: {pattern}, go next pattern")
+
+        return result_counts

src/kestrel_datasource_stixshifter/interface.py

@@ -92,6 +92,12 @@
 
 #. any in-session edit through the ``CONFIG`` command.
 
+Once you added data source profiles into ``stixshifter.yaml``, you can test the data source with command:
+
+.. code-block:: console
+
+    $ stix-shifter-diag data_source_name
+
 If you launch Kestrel in debug mode, STIX-shifter debug mode is still not
 enabled by default. To record debug level logs of STIX-shifter, create
 environment variable ``KESTREL_STIXSHIFTER_DEBUG`` with any value.

tests/test_cli.py

@@ -1,21 +1,10 @@
 import pytest
-import os
 import subprocess
 
-@pytest.fixture()
-def setup_huntflow(tmp_path):
-    profiles = """
-profiles:
-    lab101:
-        connector: stix_bundle
-        connection:
-            host: https://github.com/opencybersecurityalliance/data-bucket-kestrel/blob/main/stix-bundles/lab101.json?raw=true
-        config:
-            auth:
-                username:
-                password:
-"""
+from .utils import stixshifter_profile_lab101
 
+@pytest.fixture()
+def create_huntflow(tmp_path):
     huntflow = """
 procs = GET process FROM stixshifter://lab101
         WHERE name = 'svchost.exe'
@@ -24,27 +13,20 @@ def setup_huntflow(tmp_path):
 
     expected_result_lines = ["VARIABLE    TYPE  #(ENTITIES)  #(RECORDS)  directory*  file*  ipv4-addr*  ipv6-addr*  mac-addr*  network-traffic*  process*  user-account*  x-ecs-destination*  x-ecs-network*  x-ecs-process*  x-ecs-source*  x-ecs-user*  x-oca-asset*  x-oca-event*", "   procs process          389        1066        1078   1114        3190        1910       1066              1014       725           1062                2016            2016            2120           2024         2124          1066          2132"]
 
-    profile_file = tmp_path / "stixshifter.yaml"
     huntflow_file = tmp_path / "hunt101.hf"
 
-    os.environ["KESTREL_STIXSHIFTER_CONFIG"] = str(profile_file.expanduser().resolve())
-    with open(profile_file, "w") as pf:
-        pf.write(profiles)
-
     with open(huntflow_file, "w") as hf:
         hf.write(huntflow)
 
     huntflow_file_path = str(huntflow_file.expanduser().resolve())
 
-    # https://docs.pytest.org/en/latest/how-to/fixtures.html#teardown-cleanup-aka-fixture-finalization
-    yield huntflow_file_path, expected_result_lines
-    del os.environ["KESTREL_STIXSHIFTER_CONFIG"]
+    return huntflow_file_path, expected_result_lines
 
 
 
-def test_cli(setup_huntflow):
+def test_cli(create_huntflow, stixshifter_profile_lab101):
 
-    huntflow_file_path, expected_result_lines = setup_huntflow
+    huntflow_file_path, expected_result_lines = create_huntflow
     result = subprocess.run(args = ["kestrel", huntflow_file_path], 
                             universal_newlines = True,
                             stdout = subprocess.PIPE
@@ -55,9 +37,9 @@ def test_cli(setup_huntflow):
     assert result_lines[-2] == expected_result_lines[1]
 
 
-def test_python_module_call(setup_huntflow):
+def test_python_module_call(create_huntflow, stixshifter_profile_lab101):
 
-    huntflow_file_path, expected_result_lines = setup_huntflow
+    huntflow_file_path, expected_result_lines = create_huntflow
     result = subprocess.run(args = ["python", "-m", "kestrel", huntflow_file_path], 
                             universal_newlines = True,
                             stdout = subprocess.PIPE