From c81fbcffeab4c3516f62994a9ae99a16caaa973f Mon Sep 17 00:00:00 2001 From: Kristinn Gudjonsson Date: Fri, 4 Feb 2022 15:07:00 +0000 Subject: [PATCH 1/4] Added a simple Mitre Att&ck class as well as a helper for Alerts API. --- jupyter/docker/docker_build/Dockerfile | 2 +- jupyter/laceworkjupyter/features/__init__.py | 1 + jupyter/laceworkjupyter/features/helper.py | 54 ++++--- jupyter/laceworkjupyter/features/mitre.py | 148 ++++++++++++++++++ .../laceworkjupyter/features/query_builder.py | 2 +- jupyter/laceworkjupyter/plugins/__init__.py | 4 +- jupyter/laceworkjupyter/plugins/alerts.py | 31 ++++ 7 files changed, 216 insertions(+), 26 deletions(-) create mode 100644 jupyter/laceworkjupyter/features/mitre.py create mode 100644 jupyter/laceworkjupyter/plugins/alerts.py diff --git a/jupyter/docker/docker_build/Dockerfile b/jupyter/docker/docker_build/Dockerfile index 79e5fcf..f80c0c1 100644 --- a/jupyter/docker/docker_build/Dockerfile +++ b/jupyter/docker/docker_build/Dockerfile @@ -32,7 +32,7 @@ COPY --chown=1000:1000 docker/docker_build/lacework /home/lacework/lacenv/share/ RUN pip install --upgrade pip setuptools wheel && \ pip install --upgrade ipywidgets jupyter_contrib_nbextensions jupyter_http_over_ws ipydatetime tabulate && \ pip install --upgrade scikit-learn matplotlib python-evtx Evtx timesketch_import_client "snowflake-connector-python[secure-local-storage,pandas]" && \ - pip install --upgrade ipyaggrid keras nbformat numpy pandas pyparsing qgrid ruamel.yaml sklearn && \ + pip install --upgrade ipyaggrid keras nbformat numpy pandas pyparsing qgrid ruamel.yaml sklearn mitreattack-python && \ pip install --upgrade tensorflow tqdm traitlets xmltodict ds4n6-lib picatrix timesketch_api_client openpyxl && \ cd /home/lacework/code && pip install -e . && \ cd /home/lacework/code/jupyter && pip install -e . && \ diff --git a/jupyter/laceworkjupyter/features/__init__.py b/jupyter/laceworkjupyter/features/__init__.py index 2b31dcc..bd120f7 100644 --- a/jupyter/laceworkjupyter/features/__init__.py +++ b/jupyter/laceworkjupyter/features/__init__.py @@ -6,6 +6,7 @@ from . import date from . import helper from . import hunt +from . import mitre from . import policies from . import query from . import query_builder diff --git a/jupyter/laceworkjupyter/features/helper.py b/jupyter/laceworkjupyter/features/helper.py index d5c89f1..4d4f0a6 100644 --- a/jupyter/laceworkjupyter/features/helper.py +++ b/jupyter/laceworkjupyter/features/helper.py @@ -14,6 +14,36 @@ logger = logging.getLogger("lacework_sdk.jupyter.feature.helper") +def extract_json_field(json_obj, item): + """ + Extract a field from a JSON struct. + + :param json_obj: Either a JSON string or a dict. + :param str item: The item string, dot delimited. + :return: The extracted field. + """ + if isinstance(json_obj, str): + try: + json_obj = json.loads(json_obj) + except json.JSONDecodeError: + logger.error("Unable to decode JSON string: %s", json_obj) + return np.nan + + if not isinstance(json_obj, dict): + logger.error("Unable to extract, not a dict: %s", type(json_obj)) + return np.nan + + data = json_obj + for point in item.split("."): + if not isinstance(data, dict): + logger.error( + "Sub-item %s is not a dict (%s)", point, type(data)) + return np.nan + + data = data.get(point) + return data + + @manager.register_feature def deep_extract_field(data_frame, column, field_string, ctx=None): """ @@ -31,31 +61,9 @@ def deep_extract_field(data_frame, column, field_string, ctx=None): :param obj ctx: The context object. :return: A pandas Series with the extracted value. """ - def _extract_function(json_obj, item): - if isinstance(json_obj, str): - try: - json_obj = json.loads(json_obj) - except json.JSONDecodeError: - logger.error("Unable to decode JSON string: %s", json_obj) - return np.nan - - if not isinstance(json_obj, dict): - logger.error("Unable to extract, not a dict: %s", type(json_obj)) - return np.nan - - data = json_obj - for point in item.split("."): - if not isinstance(data, dict): - logger.error( - "Sub-item %s is not a dict (%s)", point, type(data)) - return np.nan - - data = data.get(point) - return data - if column not in data_frame: logger.error("Column does not exist in the dataframe.") return pd.Series() return data_frame[column].apply( - lambda x: _extract_function(x, field_string)) + lambda x: extract_json_field(x, field_string)) diff --git a/jupyter/laceworkjupyter/features/mitre.py b/jupyter/laceworkjupyter/features/mitre.py new file mode 100644 index 0000000..20c8820 --- /dev/null +++ b/jupyter/laceworkjupyter/features/mitre.py @@ -0,0 +1,148 @@ +""" +File that provides access to Mitre ATT&CK data. +""" + +import pandas as pd + +from taxii2client.v20 import Server +from stix2 import CompositeDataSource + +from laceworkjupyter import manager + + +class MitreAttack: + """ + Simple class that controls access to Mitre Att&ck datasets. + """ + + DEFAULT_SERVER = "https://cti-taxii.mitre.org/taxii/" + + def __init__(self): + self._api_root = None + self._server = None + self._source = None + self._frames = {} + + def _get_frame_from_collection(self, collection): + """ + Returns a DataFrame with information from a Mitre Att&ck collection. + """ + data = collection.get_objects() + lines = [] + for item in data.get('objects', []): + line = item + external = item.get('external_references', []) + for source in external: + if not source.get("url", "").startswith("https://attack.mitre.org"): + continue + line["mitre_source"] = source.get("source_name") + line["mitre_id"] = source.get("external_id") + url = source.get("url") + url_items = url.split("/") + line["mitre_collection"] = url_items[-2] + lines.append(line) + return pd.DataFrame(lines) + + @property + def server(self): + """ + Returns a copy of the Mitre server object. + """ + if self._server: + return self._server + + self._server = Server(self.DEFAULT_SERVER) + self._api_root = self._server.api_roots[0] + return self._server + + @property + def collection_titles(self): + """ + Returns a list of collection titles and ID that are available. + """ + _ = self.server + return [(c.title, c.id) for c in self._api_root.collections] + + @property + def source(self): + """ + Returns the data source object for Att&ck data. + """ + if self._source: + return self._source + + self._source = CompositeDataSource() + self._source.add_data_sources(self._api_root.collections) + + return self._source + + def get_collection(self, collection_title): + """ + Returns a DataFrame with the content of a particular collection. + + :param str collection_title: The tile or a substring of a collection + entry in Mitre Att&ck framework. The first match of available + collections will be used. + :return: A pandas DataFrame with the content or an empty frame if the + collection was not found. + """ + _ = self.server + for collection in self._api_root.collections: + if collection_title.lower() not in collection.title.lower(): + continue + + title = collection.title + if title not in self._frames: + self._frames[title] = self._get_frame_from_collection( + collection) + return self._frames[title] + + return pd.DataFrame() + + def get_technique(self, mitre_id, collection_title='Enterprise ATT&CK'): + """ + Returns a pandas Series with the content of a single technique. + + :param str mitre_id: This is the corresponding Mitre Att&ck ID for the + technique to look up. + :param str collection_title: The title of the collection this + particular technique exists in. By default uses the Enterprise + collection. + :return: A pandas Series with the information, or an empty series + if not found. + """ + titles = [x[0] for x in self.collection_titles] + use_title = "" + for title in titles: + if collection_title.lower() in title.lower(): + use_title = title + break + if not use_title: + raise ValueError("Collection Title %s not found", collection_title) + + if use_title not in self._frames: + _ = self.get_collection(use_title) + frame = self._frames.get(use_title) + + frame_slice = frame[frame.mitre_id == mitre_id.upper()] + if not frame_slice.empty: + return frame_slice.iloc[0].dropna() + + return pd.Series() + + +@manager.register_feature +def get_mitre_client(ctx=None): + """ + Returns a Mitre class object. + """ + if ctx: + mitre_client = ctx.get("mitre_client") + if mitre_client: + return mitre_client + + mitre_client = MitreAttack() + ctx.add("mitre_client", mitre_client) + return mitre_client + + return MitreAttack() diff --git a/jupyter/laceworkjupyter/features/query_builder.py b/jupyter/laceworkjupyter/features/query_builder.py index 84e9eb4..ffaf987 100644 --- a/jupyter/laceworkjupyter/features/query_builder.py +++ b/jupyter/laceworkjupyter/features/query_builder.py @@ -171,7 +171,7 @@ def generate_filters(change): # noqa: C901 table_name = datasource['name'] lw_ctx.add_state("query_builder", "query_custom_table_name", table_name) - table_schema = lw_ctx.client.datasources.get_datasource_schema(table_name) + table_schema = lw_ctx.client.datasources.get_datasource(table_name) checkboxes = [] return_fields = [] diff --git a/jupyter/laceworkjupyter/plugins/__init__.py b/jupyter/laceworkjupyter/plugins/__init__.py index 47801cc..0e9e2f9 100644 --- a/jupyter/laceworkjupyter/plugins/__init__.py +++ b/jupyter/laceworkjupyter/plugins/__init__.py @@ -1,11 +1,13 @@ """A simple loading of plugins.""" +from . import alerts from . import alert_rules from . import datasources PLUGINS = { + 'alerts.get': alerts.process_alerts, 'alert_rules.get': alert_rules.process_alert_rules, 'datasources.list_data_sources': datasources.process_list_data_sources, - 'datasources.get_datasource_schema': datasources.process_datasource_schema, + 'datasources.get_datasource': datasources.process_datasource_schema, } diff --git a/jupyter/laceworkjupyter/plugins/alerts.py b/jupyter/laceworkjupyter/plugins/alerts.py new file mode 100644 index 0000000..1640634 --- /dev/null +++ b/jupyter/laceworkjupyter/plugins/alerts.py @@ -0,0 +1,31 @@ +"""An output plugin for the Lacework Alerts API.""" + + +import pandas as pd + +from laceworkjupyter.features import helper + + +def process_alerts(data): + """ + Returns a Pandas DataFrame from the API call. + + :return: A pandas DataFrame. + """ + data_dicts = data.get("data", []) + lines = [] + for data_dict in data_dicts: + data_dict["alertDescription"] = helper.extract_json_field( + data_dict.get("alertProps", {}), "description.descriptionId") + + description_dict = helper.extract_json_field( + data_dict.get("alertProps", {}), "description.descriptionObj") + data_dict.update(description_dict) + + alert_context = helper.extract_json_field( + data_dict.get("keys", {}), "src.keys.alert") + if alert_context: + data_dict.update(alert_context) + + lines.append(data_dict) + return pd.DataFrame(lines) From 96d2f1b93a74828e3eaa49a85ab3e3a2cf4e46e5 Mon Sep 17 00:00:00 2001 From: Kristinn Gudjonsson Date: Fri, 4 Feb 2022 15:30:18 +0000 Subject: [PATCH 2/4] Added a simple YAML definition and a merge operation. --- jupyter/laceworkjupyter/features/mitre.py | 39 ++++++++++++++++++++- jupyter/laceworkjupyter/features/mitre.yaml | 5 +++ jupyter/laceworkjupyter/version.py | 2 +- jupyter/setup.py | 1 + 4 files changed, 45 insertions(+), 2 deletions(-) create mode 100644 jupyter/laceworkjupyter/features/mitre.yaml diff --git a/jupyter/laceworkjupyter/features/mitre.py b/jupyter/laceworkjupyter/features/mitre.py index 20c8820..a22e350 100644 --- a/jupyter/laceworkjupyter/features/mitre.py +++ b/jupyter/laceworkjupyter/features/mitre.py @@ -8,6 +8,9 @@ from stix2 import CompositeDataSource from laceworkjupyter import manager +from laceworkjupyter import utils as main_utils +from laceworkjupyter.features import client +from laceworkjupyter.features import utils class MitreAttack: @@ -144,5 +147,39 @@ def get_mitre_client(ctx=None): mitre_client = MitreAttack() ctx.add("mitre_client", mitre_client) return mitre_client - + return MitreAttack() + + +@manager.register_feature +def get_alerts_data_with_mitre(start_time="", end_time="", ctx=None): + """ + Returns a DataFrame from the Alerts API call with Att&ck information. + + :param str start_time: The start time, in ISO format. + :param str end_time: The end time, in ISO format. + :return: A pandas DataFrame with the content of the Alerts API call + merged with Att&CK information, if applicable. + """ + if not ctx: + raise ValueError("The context is required for this operation.") + + attack_mappings = pd.DataFrame(utils.load_yaml_file("mitre.yaml")) + mitre_client = get_mitre_client(ctx=ctx) + mitre_enterprise = mitre_client.get_collection("enterprise") + mitre_joined_df = attack_mappings.merge(mitre_enterprise, how="left") + + lw_client = ctx.client + if not lw_client: + lw_client = client.get_client() + + default_start, default_end = main_utils.parse_date_offset("LAST 2 DAYS") + if not start_time: + start_time = ctx.get("start_time", default_start) + if not end_time: + end_time = ctx.get("end_time", default_end) + + alert_df = lw_client.alerts.get(start_time=start_time, end_time=end_time) + return alert_df.merge(mitre_joined_df, how="left") + + diff --git a/jupyter/laceworkjupyter/features/mitre.yaml b/jupyter/laceworkjupyter/features/mitre.yaml new file mode 100644 index 0000000..cde5bc6 --- /dev/null +++ b/jupyter/laceworkjupyter/features/mitre.yaml @@ -0,0 +1,5 @@ +- mitre_id: T1204 + alertType: SuspiciousApplicationLaunched + +- mitre_id: T1531 + alertType: NetworkSecurityGroupCreatedOrUpdated diff --git a/jupyter/laceworkjupyter/version.py b/jupyter/laceworkjupyter/version.py index 35fabf1..ef3a7a1 100644 --- a/jupyter/laceworkjupyter/version.py +++ b/jupyter/laceworkjupyter/version.py @@ -1,4 +1,4 @@ -__version__ = '0.1.2' +__version__ = '0.1.5' def get_version(): diff --git a/jupyter/setup.py b/jupyter/setup.py index 1efb24b..0aa3065 100644 --- a/jupyter/setup.py +++ b/jupyter/setup.py @@ -35,6 +35,7 @@ "pyyaml", "laceworksdk", "ipywidgets", + "mitreattack-python", "pandas" ], classifiers=[ From c31097790d65ed20b14d3e9aa00d8b6e6d145d52 Mon Sep 17 00:00:00 2001 From: Kristinn Gudjonsson Date: Fri, 4 Feb 2022 15:35:42 +0000 Subject: [PATCH 3/4] linter --- jupyter/laceworkjupyter/features/mitre.py | 2 -- jupyter/setup.py | 1 + 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/jupyter/laceworkjupyter/features/mitre.py b/jupyter/laceworkjupyter/features/mitre.py index a22e350..c9feee0 100644 --- a/jupyter/laceworkjupyter/features/mitre.py +++ b/jupyter/laceworkjupyter/features/mitre.py @@ -181,5 +181,3 @@ def get_alerts_data_with_mitre(start_time="", end_time="", ctx=None): alert_df = lw_client.alerts.get(start_time=start_time, end_time=end_time) return alert_df.merge(mitre_joined_df, how="left") - - diff --git a/jupyter/setup.py b/jupyter/setup.py index 0aa3065..5e83179 100644 --- a/jupyter/setup.py +++ b/jupyter/setup.py @@ -36,6 +36,7 @@ "laceworksdk", "ipywidgets", "mitreattack-python", + "taxii2client", "pandas" ], classifiers=[ From ddd3df84d591c7e4caf7c1816a1a6407ad6c1f4c Mon Sep 17 00:00:00 2001 From: Kristinn Gudjonsson Date: Fri, 4 Feb 2022 15:38:03 +0000 Subject: [PATCH 4/4] adding into requirements --- jupyter/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/jupyter/requirements.txt b/jupyter/requirements.txt index 00334c7..a22fda5 100644 --- a/jupyter/requirements.txt +++ b/jupyter/requirements.txt @@ -2,3 +2,4 @@ laceworksdk>=0.9.23 pandas>=1.0.1 pyyaml>=5.4.1 ipywidgets>=7.6.5 +mitreattack-python>=1.4.2