fix: various lint fixes

Author: elucchesileon

mitreattack/attackToExcel/attackToExcel.py

@@ -2,7 +2,7 @@
 
 import argparse
 import os
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import pandas as pd
 import requests
@@ -16,7 +16,7 @@
 SUB_CHARACTERS = ["\\", "/"]
 
 
-def get_stix_data(domain: str, version: str = None, remote: str = None, stix_file: str = None) -> MemoryStore:
+def get_stix_data(domain: str, version: Optional[str] = None, remote: Optional[str] = None, stix_file: Optional[str] = None) -> MemoryStore:
     """Download the ATT&CK STIX data for the given domain and version from MITRE/CTI (or just domain if a remote workbench is specified).
 
     Parameters
@@ -110,7 +110,7 @@ def build_dataframes(src: MemoryStore, domain: str) -> Dict:
     return df
 
 
-def write_excel(dataframes: Dict, domain: str, version: str = None, output_dir: str = ".") -> List:
+def write_excel(dataframes: Dict, domain: str, 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
@@ -262,11 +262,11 @@ def write_excel(dataframes: Dict, domain: str, version: str = None, output_dir:
 
 def export(
     domain: str = "enterprise-attack",
-    version: str = None,
+    version: Optional[str] = None,
     output_dir: str = ".",
-    remote: str = None,
-    stix_file: str = None,
-    mem_store: MemoryStore = None,
+    remote: Optional[str] = None,
+    stix_file: Optional[str] = None,
+    mem_store: Optional[MemoryStore] = None,
 ):
     """Download ATT&CK data from MITRE/CTI and convert it to Excel spreadsheets.
 
@@ -298,6 +298,8 @@ def export(
     ------
     TypeError
         Raised when missing exactly one of `remote`, `stix_file`, or `mem_store`.
+    ValueError
+        Raised when `mem_store` fails to load.
     """
     if (
         (remote and stix_file and mem_store)
@@ -312,6 +314,9 @@ def export(
     if get_stix_from_github or remote or stix_file:
         mem_store = get_stix_data(domain=domain, version=version, remote=remote, stix_file=stix_file)
 
+    if mem_store is None:
+        raise ValueError("`mem_store` is empty - this should not be possible!")
+
     logger.info(f"************ Exporting {domain} to Excel ************")
 
     # build dataframes

mitreattack/attackToExcel/stixToDf.py

@@ -14,6 +14,7 @@
 from mitreattack.constants import MITRE_ATTACK_ID_SOURCE_NAMES, PLATFORMS_LOOKUP
 from mitreattack.stix20 import MitreAttackData
 
+
 def remove_revoked_deprecated(stix_objects):
     """Remove any revoked or deprecated objects from queries made to the data source."""
     # Note we use .get() because the property may not be present in the JSON data. The default is False
@@ -132,6 +133,8 @@ def techniquesToDf(src, domain):
         if subtechnique:
             subtechnique_of = all_sub_techniques.query([Filter("source_ref", "=", technique["id"])])[0]
             parent = src.get(subtechnique_of["target_ref"])
+        else:
+            parent = None
 
         # base STIX properties
         row = parseBaseStix(technique)
@@ -164,7 +167,7 @@ def techniquesToDf(src, domain):
         # domain specific fields -- enterprise
         if domain == "enterprise-attack":
             row["is sub-technique"] = subtechnique
-            if subtechnique:
+            if subtechnique and parent is not None:
                 row["name"] = f"{parent['name']}: {technique['name']}"
                 row["sub-technique of"] = parent["external_references"][0]["external_id"]
 

mitreattack/collections/collection_to_index.py

@@ -75,12 +75,12 @@ def generate_index(name, description, root_url, files=None, folders=None, sets=N
                         )
                         cleaned_bundles.append(potentially_valid_bundle)
                     else:
-                        print(f"cannot use bundle {potentially_valid_bundle.id} due to lack of collection object")
+                        print(f"cannot use bundle {potentially_valid_bundle['id']} due to lack of collection object")
                 else:
-                    print(f"cannot use bundle {potentially_valid_bundle.id} due to lack of collection object")
+                    print(f"cannot use bundle {potentially_valid_bundle['id']} due to lack of collection object")
 
-        index_created = None
-        index_modified = None
+        index_created = isoparse("9999-12-31T23:59:59.999Z")
+        index_modified = isoparse("9999-12-31T23:59:59.999Z")
         collections = {}  # STIX ID -> collection object
 
         if files:

mitreattack/download_stix.py

@@ -27,6 +27,11 @@ def download_stix(stix_version: str, domain: str, download_dir: str, release: st
         ATT&CK release to download.
     known_hash : str
         SHA256 hash of the ATT&CK release.
+
+    Raises
+    ------
+    ValueError
+        Raised if `stix_version` is not "2.0" or "2.1".
     """
     release_download_dir = pathlib.Path(f"{download_dir}/v{release}")
     release_download_dir.mkdir(parents=True, exist_ok=True)
@@ -36,6 +41,8 @@ def download_stix(stix_version: str, domain: str, download_dir: str, release: st
         download_url = f"https://raw.githubusercontent.com/mitre/cti/ATT%26CK-v{release}/{domain}-attack/{fname}"
     elif stix_version == "2.1":
         download_url = f"https://raw.githubusercontent.com/mitre-attack/attack-stix-data/master/{domain}-attack/{domain}-attack-{release}.json"
+    else:
+        raise ValueError(f"Invalid STIX version: {stix_version}")
 
     pooch.retrieve(download_url, known_hash=known_hash, fname=fname, path=str(release_download_dir))
 
@@ -57,6 +64,11 @@ def download_domains(
         Version of STIX to download. Options are "2.0" or "2.1"
     attack_versions : List[str], optional
         List of specific ATT&CK versions to download. If provided, overrides all_versions behavior.
+
+    Raises
+    ------
+    ValueError
+        Raised if `stix_version` is not "2.0" or "2.1".
     """
     for domain in domains:
         if domain == "pre" and stix_version == "2.1":
@@ -67,6 +79,8 @@ def download_domains(
             stix_hash_data = release_info.STIX20
         elif stix_version == "2.1":
             stix_hash_data = release_info.STIX21
+        else:
+            raise ValueError(f"Invalid STIX version: {stix_version}")
 
         releases = {}
         if domain == "enterprise":

mitreattack/navlayers/core/layer.py

@@ -3,7 +3,7 @@
 import json
 import traceback
 
-from mitreattack.navlayers.core.exceptions import UninitializedLayer, BadType, BadInput, handler
+from mitreattack.navlayers.core.exceptions import BadInput, BadType, UninitializedLayer, handler
 from mitreattack.navlayers.core.layerobj import _LayerObj
 
 
@@ -61,6 +61,7 @@ def from_file(self, filename):
         :param filename: the target filename to load from
         """
         fallback = False
+        raw = ""
         with open(filename, "r", encoding="utf-16") as fio:
             try:
                 raw = fio.read()

mitreattack/navlayers/exporters/excel_templates.py

@@ -1,7 +1,7 @@
 """Contains ExcelTemplates class."""
 
 import openpyxl
-from openpyxl.styles import Font, Alignment, Border, Side, PatternFill
+from openpyxl.styles import Alignment, Border, Font, PatternFill, Side
 
 from mitreattack.navlayers.exporters.matrix_gen import MatrixGen
 

mitreattack/navlayers/exporters/svg_objects.py

@@ -1,9 +1,10 @@
 """Helper classes and functions for working with SVG objects."""
 
-import drawsvg
 import colorsys
-import numpy as np
 import os
+
+import drawsvg
+import numpy as np
 from PIL import ImageFont
 
 from mitreattack.navlayers.core.gradient import Gradient