add mitre-attack

fix errors when loading library without framework
intuitem · Nov 22, 2023 · ea87938 · ea87938
1 parent 4d62f86
commit ea87938
Show file tree

Hide file tree

Showing 4 changed files with 5,269 additions and 16 deletions.
diff --git a/core/helpers.py b/core/helpers.py
@@ -230,4 +230,4 @@ def get_sorted_requirements_and_groups_rec(requirements: list, requirement_group
         requirements,
         requirement_groups,
         requirements_assessed,
-        [rg for rg in requirement_groups if not rg.parent_urn])
+        [rg for rg in requirement_groups if not rg.parent_urn] if requirement_groups else [])
diff --git a/library/helpers.py b/library/helpers.py
@@ -8,23 +8,24 @@ def preview_library(library) -> list:
     preview = {}
     requirement_groups_list = []
     requirements_list = []
-    if library['objects']['framework'].get('requirement_groups'):
-        for requirement_group in library['objects']['framework']['requirement_groups']:
+    framework = library['objects'].get('framework')
+    if framework:
+        for requirement_group in library['objects']['framework'].get('requirement_groups'):
             requirement_groups_list.append(
                 RequirementGroup(description = requirement_group.get('description'),
                                 name = requirement_group['name'],
                                 urn = requirement_group['urn'],
                                 parent_urn = requirement_group.get('parent_urn')))
-    for requirement in library['objects']['framework']['requirements']:
-        if Requirement.objects.filter(urn=requirement['urn']).exists():
-            requirements_list.append(Requirement.objects.get(urn=requirement['urn']))
-        else:
-            temp_req = Requirement(description = requirement.get('description'),
-                                                    name = requirement['name'],
-                                                    urn = requirement['urn'],
-                                                    parent_urn = requirement.get('parent_urn'),
-                                                    folder = Folder.get_root_folder())
-            requirements_list.append(temp_req)
-    preview['requirement_groups'] = requirement_groups_list
-    preview['requirements'] = requirements_list
+        for requirement in library['objects']['framework'].get('requirements'):
+            if Requirement.objects.filter(urn=requirement['urn']).exists():
+                requirements_list.append(Requirement.objects.get(urn=requirement['urn']))
+            else:
+                temp_req = Requirement(description = requirement.get('description'),
+                                                        name = requirement['name'],
+                                                        urn = requirement['urn'],
+                                                        parent_urn = requirement.get('parent_urn'),
+                                                        folder = Folder.get_root_folder())
+                requirements_list.append(temp_req)
+        preview['requirement_groups'] = requirement_groups_list
+        preview['requirements'] = requirements_list
     return preview