Merge pull request #146 from arkid15r/ark/86-increase-test-coverage

Increase test coverage

.github/workflows/lint-publish.yml

@@ -22,11 +22,10 @@ jobs:
       - name: Lint with flake8
         run: flake8 mitreattack/ --count --exit-zero --statistics
 
-      # should turn these back on once they take less than 10 minutes to run
-      # - name: Run pytest
-      #   run: |
-      #     cd tests
-      #     pytest --cov=mitreattack --cov-report html
+      - name: Run pytest
+        run: |
+          cd tests
+          pytest --cov=mitreattack --cov-report html
 
   publish:
     needs: lint-test

docs/conf.py

@@ -4,28 +4,28 @@
 
 import os
 import sys
-sys.path.insert(0, os.path.abspath('..'))
+
+sys.path.insert(0, os.path.abspath(".."))
 
 # -- Project information -----------------------------------------------------
-project = 'mitreattack-python'
-copyright = '2022, The MITRE Corporation'
-version = '2.0.0'
-release = '2.0.0'
+project = "mitreattack-python"
+copyright = "2022, The MITRE Corporation"
+version = "2.0.0"
+release = "2.0.0"
 
 # -- General configuration ---------------------------------------------------
 extensions = [
-    'sphinx.ext.autodoc',
-    'sphinx.ext.autosummary',
-    'sphinx.ext.autosectionlabel',
-    'sphinx.ext.napoleon'
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.autosectionlabel",
+    "sphinx.ext.napoleon",
 ]
 autosummary_generate = True
 autodoc_default_options = {"members": True}
 add_module_names = True
-templates_path = ['_templates']
-exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 # -- Options for HTML output -------------------------------------------------
-html_theme = 'sphinx_rtd_theme'
-html_static_path = ['_static']
-
+html_theme = "sphinx_rtd_theme"
+html_static_path = ["_static"]

examples/get_procedure_examples_by_tactic.py

@@ -30,12 +30,22 @@ def main():
         for technique in techniques:
             technique_stix_id = technique["id"]
             groups_using_technique = mitre_attack_data.get_groups_using_technique(technique_stix_id=technique_stix_id)
-            software_using_technique = mitre_attack_data.get_software_using_technique(technique_stix_id=technique_stix_id)
-            campaigns_using_technique = mitre_attack_data.get_campaigns_using_technique(technique_stix_id=technique_stix_id)
-
-            print_procedure_examples(mitre_attack_data=mitre_attack_data, attack_objects_using_technique=groups_using_technique)
-            print_procedure_examples(mitre_attack_data=mitre_attack_data, attack_objects_using_technique=software_using_technique)
-            print_procedure_examples(mitre_attack_data=mitre_attack_data, attack_objects_using_technique=campaigns_using_technique)
+            software_using_technique = mitre_attack_data.get_software_using_technique(
+                technique_stix_id=technique_stix_id
+            )
+            campaigns_using_technique = mitre_attack_data.get_campaigns_using_technique(
+                technique_stix_id=technique_stix_id
+            )
+
+            print_procedure_examples(
+                mitre_attack_data=mitre_attack_data, attack_objects_using_technique=groups_using_technique
+            )
+            print_procedure_examples(
+                mitre_attack_data=mitre_attack_data, attack_objects_using_technique=software_using_technique
+            )
+            print_procedure_examples(
+                mitre_attack_data=mitre_attack_data, attack_objects_using_technique=campaigns_using_technique
+            )
 
 
 if __name__ == "__main__":

examples/get_revoked_techniques.py

@@ -14,7 +14,9 @@ def main():
             revoking_object = mitre_attack_data.get_revoking_object(revoked_stix_id=stix_id)
             revoking_object_name = revoking_object.get("name")
             revoking_object_attack_id = mitre_attack_data.get_attack_id(stix_id=revoking_object["id"])
-            print(f"[{revoked_technique_attack_id}] {revoked_technique_name} was revoked by [{revoking_object_attack_id}] {revoking_object_name}")
+            print(
+                f"[{revoked_technique_attack_id}] {revoked_technique_name} was revoked by [{revoking_object_attack_id}] {revoking_object_name}"
+            )
 
 
 if __name__ == "__main__":

mitreattack/__init__.py

@@ -1,3 +1,3 @@
 from .attackToExcel import *
 from .navlayers import *
-from .collections import *
+from .collections import *

mitreattack/collections/collection_to_index.py

@@ -171,7 +171,7 @@ def main(args):
             name=args.name, description=args.description, root_url=args.root_url, files=args.files, folders=args.folders
         )
         print(f"writing {args.output}")
-        json.dump(index, f, indent=4, enforce_ascii=False)
+        json.dump(index, f, indent=4, ensure_ascii=False)
 
 
 if __name__ == "__main__":

mitreattack/collections/stix_to_collection.py

@@ -30,7 +30,9 @@ def stix_to_collection(bundle, name, version, description=None):
         for obj in working_bundle["objects"]:  # check to see if this bundle already contains a collection
             if obj["type"] == "x-mitre-collection":
                 return bundle
-        if bundle.get("spec_version", "") == "2.0":
+
+        bundle_version = bundle.get("spec_version", "")
+        if bundle_version == "2.0":
             try:
                 print(
                     "[NOTE] - version 2.0 spec detected. Forcibly upgrading the bundle to 2.1 to support "
@@ -51,11 +53,13 @@ def stix_to_collection(bundle, name, version, description=None):
                 )
                 print(f"[ERROR] - Full Error trace: {traceback.print_exc(e)}")
                 return None
-        if bundle.get("spec_version", "") != "2.1" and bundle.get("spec_version", "") != "2.0":
+        elif bundle_version != "2.1":
             print(
-                f"[ERROR] - version {working_bundle.get('spec_version', '[NOT FOUND]')} is not one of [2.0, 2.1]. "
+                f"[ERROR] - version {bundle_version or '[NOT FOUND]'} is not one of [2.0, 2.1]. "
                 f"This module only processes stix 2.0 and stix 2.1 bundles."
             )
+            return None
+
         time = datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%fZ")
         if not description:
             description = "This collection was autogenerated by STIXToCollection, as part of mitreattack-python"
@@ -103,14 +107,20 @@ def stix_to_collection(bundle, name, version, description=None):
 
 def main(args):
     """Entrypoint for stixToCollection_cli."""
-    with open(args.input, "r", encoding="utf-16") as f:
-        bundle = json.load(f)
-        with open(args.output, "w", encoding="utf-16") as f2:
-            f2.write(
-                json.dumps(
-                    STIXToCollection.stix_to_collection(bundle, args.name, args.version, args.description), indent=4
-                )
+    with open(args.input, "r", encoding="utf-16") as input:
+        bundle = json.load(input)
+    with open(args.output, "w", encoding="utf-16") as output:
+        output.write(
+            json.dumps(
+                STIXToCollection.stix_to_collection(
+                    bundle,
+                    args.name,
+                    args.version,
+                    args.description,
+                ),
+                indent=4,
             )
+        )
 
 
 if __name__ == "__main__":

mitreattack/stix20/__init__.py

@@ -1,2 +1,2 @@
 from .MitreAttackData import MitreAttackData
-from .custom_attack_objects import StixObjectFactory, Matrix, Tactic, DataSource, DataComponent
+from .custom_attack_objects import StixObjectFactory, Matrix, Tactic, DataSource, DataComponent

tests/conftest.py

@@ -19,7 +19,7 @@ def attack_stix_dir():
         domains=["enterprise", "mobile", "ics"],
         download_dir=download_dir,
         all_versions=False,
-        stix_version="2.0"
+        stix_version="2.0",
     )
 
     yield download_dir
@@ -29,20 +29,17 @@ def attack_stix_dir():
 
 @pytest.fixture(scope="session")
 def stix_file_enterprise_latest(attack_stix_dir):
-    stix_file = f"{attack_stix_dir}/v{LATEST_VERSION}/enterprise-attack.json"
-    return stix_file
+    return f"{attack_stix_dir}/v{LATEST_VERSION}/enterprise-attack.json"
 
 
 @pytest.fixture(scope="session")
 def stix_file_mobile_latest(attack_stix_dir):
-    stix_file = f"{attack_stix_dir}/v{LATEST_VERSION}/mobile-attack.json"
-    return stix_file
+    return f"{attack_stix_dir}/v{LATEST_VERSION}/mobile-attack.json"
 
 
 @pytest.fixture(scope="session")
 def stix_file_ics_latest(attack_stix_dir):
-    stix_file = f"{attack_stix_dir}/v{LATEST_VERSION}/ics-attack.json"
-    return stix_file
+    return f"{attack_stix_dir}/v{LATEST_VERSION}/ics-attack.json"
 
 
 @pytest.fixture(scope="session")