Efficency pass 1 - now using objects to clean things up a bit

Author: clittle

layers/exporters/excel_templates.py

@@ -52,11 +52,11 @@ def _build_raw(self, showName=True, showID=False, subtechs=[], exclude=[]):
             c = sheet.cell(row=entry[0], column=entry[1])
             write_val = ''
             if showName and showID:
-                write_val = self._get_ID(template[entry]) + ': ' + template[entry]
+                write_val = MatrixGen._get_ID(self.codex, template[entry]) + ': ' + template[entry]
             elif showName:
                 write_val = template[entry]
             elif showID:
-                write_val = self._get_ID(template[entry])
+                write_val = MatrixGen._get_ID(self.codex, template[entry])
             c.value = write_val
             if entry[0] == 1:
                 c.font = header_template_f
@@ -83,41 +83,6 @@ def _build_raw(self, showName=True, showID=False, subtechs=[], exclude=[]):
 
         return wb
 
-    def _get_ID(self, name):
-        """
-            INTERNAL - Do lookups to retrieve the ID of a technique given it's name
-
-            :param name: The name of the technique to retrieve the ID of
-            :return: The ID of the technique referenced by name
-        """
-        t_lookup = {'Initial Access': 'TA0001',
-                    'Execution': 'TA0002',
-                    'Persistence': 'TA0003',
-                    'Privilege Escalation': 'TA0004',
-                    'Defense Evasion': 'TA0005',
-                    'Credential Access': 'TA0006',
-                    'Discovery': 'TA0007',
-                    'Lateral Movement': 'TA0008',
-                    'Collection': 'TA0009',
-                    'Command and Control': 'TA0011',
-                    'Exfiltration': 'TA0010',
-                    'Impact': 'TA0040'}
-        for col in self.codex:
-            for elm in col:
-                for key in elm:
-                    if key == 'tactic':
-                        if elm[key] == name:
-                            return t_lookup[elm[key]]
-                    elif key == 'subtechs':
-                        for en in elm[key]:
-                            for ent in elm[key][en]:
-                                for sub in ent:
-                                    if name == sub:
-                                        return str(ent[sub])
-                    else:
-                        if name == key:
-                            return str(elm[key])
-
     def export(self, showName, showID, subtechs=[], exclude=[]):
         """
             Export a raw customized excel template
@@ -140,18 +105,7 @@ def retrieve_coords(self, techniqueID, tactic=None):
             :return: A tuple representing the (row, column) of the target element in the workbook
         """
         listing = []
-        match = ''
-        for col in self.codex:
-            for elm in col:
-                for name in elm:
-                    if name == "subtechs":
-                        for sp in elm[name]:
-                            for t in elm[name][sp]:
-                                for na in t:
-                                    if t[na] == techniqueID:
-                                        match = na
-                    if elm[name] == techniqueID:
-                        match = name
+        match = MatrixGen._get_name(self.codex, techniqueID)
         for entry in self.template:
             if self.template[entry] == match:
                 if tactic is not None:

layers/exporters/matrix_gen.py

@@ -2,6 +2,30 @@
 from stix2 import TAXIICollectionSource, Filter, FileSystemSource
 from taxii2client import Server, Collection
 
+class MatrixEntry:
+    def __init__(self, id=None, name=None):
+        if id is not None:
+            self.id = id
+        if name is not None:
+            self.name = name
+
+    @property
+    def id(self):
+        if self.__id is not None:
+            return self.__id
+
+    @id.setter
+    def id(self, id):
+        self.__id = id
+
+    @property
+    def name(self):
+        if self.__name is not None:
+            return self.__name
+
+    @name.setter
+    def name(self, name):
+        self.__name = name
 
 class MatrixGen:
     convert = {
@@ -18,6 +42,18 @@ class MatrixGen:
              'exfiltration': 'Exfiltration',
              'impact': 'Impact'
         }
+    t_lookup = {'Initial Access': 'TA0001',
+                'Execution': 'TA0002',
+                'Persistence': 'TA0003',
+                'Privilege Escalation': 'TA0004',
+                'Defense Evasion': 'TA0005',
+                'Credential Access': 'TA0006',
+                'Discovery': 'TA0007',
+                'Lateral Movement': 'TA0008',
+                'Collection': 'TA0009',
+                'Command and Control': 'TA0011',
+                'Exfiltration': 'TA0010',
+                'Impact': 'TA0040'}
     def __init__(self, fresh=True):
         """
             Initialization - Creates a matrix generator object
@@ -64,19 +100,19 @@ def _get_technique_listing(self, tactic, mode='enterprise'):
             :param tactic: The tactic to grab techniques from
             :param mode: The domain to draw from
         """
-        techniques = {}
+        techniques = []
         subtechs = {}
         techs = self.collections[mode].query([Filter('type', '=', 'attack-pattern'), Filter('kill_chain_phases.phase_name', '=', tactic)])
         for entry in techs:
             if entry['kill_chain_phases'][0]['kill_chain_name'] == 'mitre-attack':
                 tid = [t['external_id'] for t in entry['external_references'] if t['source_name'] == 'mitre-attack']
                 if '.' not in tid[0]:
-                    techniques[entry['name']] = tid[0]
+                    techniques.append(MatrixEntry(id=tid[0], name=entry['name']))
                 else:
                     parent = tid[0].split('.')[0]
                     if parent not in subtechs:
                         subtechs[parent] = []
-                    subtechs[parent].append({entry['name']: tid[0]})
+                    subtechs[parent].append(MatrixEntry(id=tid[0], name=entry['name']))
         return techniques, subtechs
 
     @staticmethod
@@ -107,18 +143,18 @@ def _construct_panop(codex, subtechs, excludes):
                     sr = entry[0]
                     joins.append([entry[0], column-1, len(stechs[entry[1]])])
                     for element in stechs[entry[1]]:
-                        matrix_obj[(sr, column)] = [x for x in element.keys()][0]
+                        matrix_obj[(sr, column)] = element.name
                         sr += 1
                 cycle = False
                 column += 1
             row = 2
-            matrix_obj[(1, column)] = col[0]['tactic']
-            c_name = col[0]['tactic']
+            matrix_obj[(1, column)] = col[0].name
+            c_name = col[0].name
             stechs = col[1]['subtechs']
             to_add = []
             for element in col[2:]:
-                elname = [x for x in element.keys()][0]
-                tid = element[[x for x in element.keys()][0]]
+                elname = element.name
+                tid = element.id
                 skip = False
                 for entry in range(0, len(et)):
                     if et[entry] == tid and (e_tacs[entry] == False or MatrixGen.convert[e_tacs[entry]] == c_name):
@@ -139,6 +175,48 @@ def _construct_panop(codex, subtechs, excludes):
                         row += 1
         return matrix_obj, joins
 
+    @staticmethod
+    def _get_ID(codex, name):
+        """
+            INTERNAL - Do lookups to retrieve the ID of a technique given it's name
+
+            :param codex: The list of lists matrix object (output of build_matrix)
+            :param name: The name of the technique to retrieve the ID of
+            :return: The ID of the technique referenced by name
+        """
+        for col in codex:
+            for elm in col:
+                if col.index(elm) == 1:
+                    for sp in elm['subtechs']:
+                        for t in elm['subtechs'][sp]:
+                            if t.name == name:
+                                return t.id
+                else:
+                    if elm.name == name:
+                        return elm.id
+        return ''
+
+    @staticmethod
+    def _get_name(codex, id):
+        """
+            INTERNAL - Do lookups to retrieve the name of a technique given it's ID
+
+            :param codex: The list of lists matrix object (output of build_matrix)
+            :param id: The ID of the technique to retrieve the name of
+            :return: The name of the technique referenced by id
+        """
+        for col in codex:
+            for elm in col:
+                if col.index(elm) == 1:
+                    for sp in elm['subtechs']:
+                        for t in elm['subtechs'][sp]:
+                            if t.id == id:
+                                return t.name
+                else:
+                    if elm.id == id:
+                        return elm.name
+        return ''
+
     def build_matrix(self, mode='enterprise'):
         """
             Build a ATT&CK matrix object, as a list of lists containing technique dictionaries
@@ -149,29 +227,15 @@ def build_matrix(self, mode='enterprise'):
         tacs = self._get_tactic_listing(mode)
         for tac in tacs:
             techs, subtechs = self._get_technique_listing(tac.lower().replace(' ', '-'), mode)
-            colm = [dict(tactic=tac)]
-            temp = []
+            colm = [MatrixEntry(id=self.t_lookup[tac], name=tac)]
             stemp = {}
             # sort subtechniques via id, append to column
             for par in subtechs:
-                stemp[par] = []
-                temp = []
-                for entry in subtechs[par]:
-                    for obj in entry:
-                        temp.append(entry[obj])
-                for entry in sorted(temp):
-                    obj = {}
-                    for k in subtechs[par]:
-                        for j in k:
-                            if k[j] == entry:
-                                obj[j] = entry
-                    stemp[par].append(obj)
+                subtechs[par].sort(key=lambda x: x.name)
+                stemp[par] = subtechs[par]
             colm.append({'subtechs':stemp})
-            temp = []
             # sort techniques alphabetically, append to column
-            for entry in techs:
-                temp.append(entry)
-            for entry in sorted(temp):
-                colm.append({entry:techs[entry]})
+            techs.sort(key=lambda x: x.name)
+            colm.extend(techs)
             matrix.append(colm)
         return matrix