chore: run ruff

Author: adpare

mitreattack/attackToExcel/attackToExcel.py

@@ -149,14 +149,9 @@ def build_dataframes(src: MemoryStore, domain: str) -> Dict:
 
 def build_ds_an_lg_relationships(dataframes: Dict) -> Dict[str, pd.DataFrame]:
     """Build detection-mappings.xlsx with a single DS → Analytic → LogSource sheet."""
+    ds_an = dataframes["detectionstrategies"].get("detectionstrategies-analytic", pd.DataFrame())
 
-    ds_an = dataframes["detectionstrategies"].get(
-        "detectionstrategies-analytic", pd.DataFrame()
-    )
-
-    an_ls = dataframes["analytics"].get(
-        "analytic-logsource", pd.DataFrame()
-    )
+    an_ls = dataframes["analytics"].get("analytic-logsource", pd.DataFrame())
 
     if ds_an.empty or an_ls.empty:
         combined = pd.DataFrame()
@@ -167,13 +162,12 @@ def build_ds_an_lg_relationships(dataframes: Dict) -> Dict[str, pd.DataFrame]:
             how="left",
         )
 
-    return {
-        "ds_an_ls": combined
-    }
-
+    return {"ds_an_ls": combined}
 
 
-def write_excel(dataframes: Dict, domain: str, src: MemoryStore, version: Optional[str] = None, output_dir: str = ".") -> List:
+def write_excel(
+    dataframes: Dict, domain: str, src: MemoryStore, version: Optional[str] = None, output_dir: str = "."
+) -> List:
     """Given a set of dataframes from build_dataframes, write the ATT&CK dataset to output directory.
 
     Parameters
@@ -232,10 +226,14 @@ def write_excel(dataframes: Dict, domain: str, src: MemoryStore, version: Option
                     for sheet_name in object_data:
                         logger.debug(f"Writing sheet to {fp}: {sheet_name}")
                         object_data[sheet_name].to_excel(object_writer, sheet_name=sheet_name, index=False)
-   
+
                     # Write Detection strategy - Analytics - Log sources file
-                    if object_type in add_ds_an_ls_to and isinstance(ds_an_ls_df, pd.DataFrame) and not ds_an_ls_df.empty:
-                        ds_an_ls_df.to_excel(object_writer, sheet_name="detection mappings", index=False)
+                    if (
+                        object_type in add_ds_an_ls_to
+                        and isinstance(ds_an_ls_df, pd.DataFrame)
+                        and not ds_an_ls_df.empty
+                    ):
+                        ds_an_ls_df.to_excel(object_writer, sheet_name="defensive mappings", index=False)
                 written_files.append(fp)
 
                 # add citations to master citations list
@@ -323,7 +321,7 @@ def write_excel(dataframes: Dict, domain: str, src: MemoryStore, version: Option
                 written_files.append(fp)
 
         if isinstance(ds_an_ls_df, pd.DataFrame) and not ds_an_ls_df.empty:
-            ds_an_ls_df.to_excel(master_writer, sheet_name="detection mappings", index=False)
+            ds_an_ls_df.to_excel(master_writer, sheet_name="defensive mappings", index=False)
         # remove duplicate citations and add sheet to master file
         logger.debug(f"Writing sheet to {master_fp}: citations")
         citations.drop_duplicates(subset="reference", ignore_index=True).sort_values("reference").to_excel(
@@ -404,7 +402,7 @@ def export(
             major_version = int(match.group(1))
             if major_version < 18:
                 dataframes = build_dataframes_pre_v18(src=mem_store, domain=domain)
-                write_excel(dataframes=dataframes, domain=domain, version=version, output_dir=output_dir)
+                write_excel(dataframes=dataframes, domain=domain, src=mem_store, version=version, output_dir=output_dir)
                 return
 
     dataframes = build_dataframes(src=mem_store, domain=domain)

mitreattack/attackToExcel/stixToDf.py

@@ -382,7 +382,11 @@ def analyticsToDf(src):
         for ds in detection_strategies:
             for analytic_id in ds.get("x_mitre_analytic_refs", []):
                 analytic_to_ds_map.setdefault(analytic_id, []).append(
-                    {"detection_strategy_attack_id": ds["external_references"][0]["external_id"], "detection_strategy_id": ds["id"], "detection_strategy_name": ds.get("name", "")}
+                    {
+                        "detection_strategy_attack_id": ds["external_references"][0]["external_id"],
+                        "detection_strategy_id": ds["id"],
+                        "detection_strategy_name": ds.get("name", ""),
+                    }
                 )
 
         for analytic in tqdm(analytics, desc="parsing analytics"):
@@ -404,7 +408,7 @@ def analyticsToDf(src):
                         "data_component_attack_id": data_comp_attack_id,
                         "log_source_name": logsrc.get("name", ""),
                         "channel": logsrc.get("channel", ""),
-                        "platforms": ", ".join(sorted(analytic.get("x_mitre_platforms", [])))
+                        "platforms": ", ".join(sorted(analytic.get("x_mitre_platforms", []))),
                     }
                 )
 
@@ -417,8 +421,7 @@ def analyticsToDf(src):
                         "detection_strategy_id": ds_info["detection_strategy_id"],
                         "detection_strategy_attack_id": ds_info["detection_strategy_attack_id"],
                         "detection_strategy_name": ds_info["detection_strategy_name"],
-                        "platforms": ", ".join(sorted(analytic.get("x_mitre_platforms", [])))
-
+                        "platforms": ", ".join(sorted(analytic.get("x_mitre_platforms", []))),
                     }
                 )
 
@@ -463,8 +466,7 @@ def detectionstrategiesToDf(src):
                         "detection_strategy_name": detection_strategy.get("name", ""),
                         "analytic_id": analytic_id,
                         "analytic_name": analytic_obj["external_references"][0]["external_id"],
-                        "platforms": ", ".join(sorted(analytic_obj.get("x_mitre_platforms", [])))
-
+                        "platforms": ", ".join(sorted(analytic_obj.get("x_mitre_platforms", []))),
                     }
                 )
 
@@ -525,6 +527,7 @@ def softwareToDf(src):
 
     return dataframes
 
+
 def detectionStrategiesAnalyticsLogSourcesDf(src):
     """Build a single DS -> LogSource -> Analytic dataframe directly from STIX."""
     detection_strategies = src.query([Filter("type", "=", "x-mitre-detection-strategy")])
@@ -550,22 +553,25 @@ def detectionStrategiesAnalyticsLogSourcesDf(src):
                 data_comp_id = logsrc.get("x_mitre_data_component_ref", "")
                 data_comp = src.get(data_comp_id)
 
-                rows.append({
-                    "detection_strategy_attack_id": ds_attack_id,
-                    "detection_strategy_id": ds_id,
-                    "detection_strategy_name": ds_name,
-                    "analytic_id": analytic_id,
-                    "analytic_name": analytic_attack_id,
-                    "platforms": platforms,
-                    "log_source_name": logsrc.get("name", ""),
-                    "channel": logsrc.get("channel", ""),
-                    "data_component_id": data_comp_id,
-                    "data_component_name": (data_comp.get("name", "") if data_comp else ""),
-                    "data_component_attack_id": data_comp["external_references"][0]["external_id"]
-                })
+                rows.append(
+                    {
+                        "detection_strategy_attack_id": ds_attack_id,
+                        "detection_strategy_id": ds_id,
+                        "detection_strategy_name": ds_name,
+                        "analytic_id": analytic_id,
+                        "analytic_name": analytic_attack_id,
+                        "platforms": platforms,
+                        "log_source_name": logsrc.get("name", ""),
+                        "channel": logsrc.get("channel", ""),
+                        "data_component_id": data_comp_id,
+                        "data_component_name": (data_comp.get("name", "") if data_comp else ""),
+                        "data_component_attack_id": data_comp["external_references"][0]["external_id"],
+                    }
+                )
 
     return pd.DataFrame(rows)
 
+
 def groupsToDf(src):
     """Parse STIX groups from the given data and return corresponding pandas dataframes.
 
@@ -1309,4 +1315,4 @@ def _get_relationship_citations(object_dataframe, relationship_df):
         else:
             for i in range(0, len(new_citations)):
                 new_citations[i] = ",".join([new_citations[i], subset[i]])
-    return new_citations
+    return new_citations