Updates to streamline things, clean up approach

Author: clittle

layers/README.md

@@ -21,9 +21,12 @@ This folder contains modules and scripts for working with ATT&CK Navigator layer
 #### Exporter Scripts
 | script | description |
 |:-------|:------------|
-| [matrix_gen](exporters/matrix_gen.py) | Provides a means by which to generate a matrix from raw data, either cached or from the ATT&CK TAXII server. |
-| [excel_templates](exporters/excel_templates.py) | Provides a means by which to convert a matrix into a base excel template. |
 | [to_excel](exporters/to_excel.py) | Provides a means by which to export an ATT&CK Layer to an excel file. A further breakdown can be found in the corresponding section below. |
+##### Utility Modules
+| script | description |
+|:-------|:------------|
+| [excel_templates](exporters/excel_templates.py) | Provides a means by which to convert a matrix into a base excel template. |
+| [matrix_gen](exporters/matrix_gen.py) | Provides a means by which to generate a matrix from raw data, either cached or from the ATT&CK TAXII server. |
 
 ## Layer
 The Layer class provides format validation and read/write capabilities to aid in working with ATT&CK Navigator Layers in python. It is the primary interface through which other Layer-related classes defined in the core module should be used. The Layer class API and a usage example are below.
@@ -128,35 +131,50 @@ out_layer6.to_file("C:\demo_layer6.json")                     # Save combined co
 ```
 
 ## to_excel.py
-to_excel.py provides the ToExcel class, which is a way to export an existing layer file as an Excel spreadsheet. The ToExcel class has an optional parameter for the export function, `to_file()`, that tells the exporter to build the output matrix based on live data from cti-taxii.mitre.org. Otherwise, it will use the cached subtechniques data. If matrix template lacks elements included in a layer's technique listing, those elements will not be visible in the output file.
+to_excel.py provides the ToExcel class, which is a way to export an existing layer file as an Excel 
+spreadsheet. The ToExcel class has an optional parameter for the initialization function, that 
+tells the exporter to build the output matrix based on live data from cti-taxii.mitre.org. It 
+can also use cached subtechniques data. If neither of these is used, it will default to pulling from the 
+ATT&CK cti repo. If matrix template lacks elements included in a layer's technique listing, those 
+elements will not be visible in the output file.
 
 ##### ToExcel()
 ```python
-x = ToExcel(layer=None)
+x = ToExcel(domain='enterprise', server=False, local=None)
 ```
-The ToExcel constructor takes an options layer argument during instantiation. If no layer is loaded at this time, it must be loaded later before exporting `x.layer = Layer()`.
+The ToExcel constructor takes domain, server, and local arguments during instantiation. The domain can 
+be either 'enterprise' or 'mobile', and can be pulled from a layer file as `layer.domain`. The server 
+argument tells the matrix generation tool to use the `cti-taxii` server when building the matrix. The 
+local argument tells the matrix generation tool the path to a local stix repository to use when 
+constructing the matrix. 
 
 ##### .to_file() Method
 ```python
-x.to_file(filepath="layer.xlsx", fresh=False)
+x.to_xlsx(layer=layer, filepath="layer.xlsx")
 ```
-The to_file method exports the layer file loaded in the ToExcel class (`x`), as an excel file to the 
-filepath specified. The fresh parameter is an optional parameter that bases the export on a matrix built 
-from the live data on the taxii server (`True`), or based on cached data (`False`).
-
+The to_xlsx method exports the layer file referenced as `layer`, as an excel file to the 
+`filepath` specified. 
 
 #### Example Usage
 ```python
 from layers import Layer
 from layers import ToExcel
 
+
+lay3 = Layer()
+lay3.from_file('layers/heatmap_example.json')
+exporter = ToExcel(domain=lay3.layer.domain, server=False, local=None)
+exporter.to_xlsx(layer=lay3, filepath="layer.xlsx")
+
 lay = Layer()
 lay.from_file("path/to/layer/file.json")
 # Using cached data for template
-t = ToExcel(lay)
-t.to_file("demo.xlsx")
+t = ToExcel(domain=lay.layer.domain, server=False, local=None)
+t.to_xlsx(layer=lay, filepath="demo.xlsx")
 # Using live data for templates
-t2 = ToExcel()
-t2.layer = lay
-t.to_file("demo2.xlsx", fresh=True)
+t2 = ToExcel(domain='enterprise', server=True, local=None)
+t.to_xlsx(layer=lay, filepath="demo2.xlsx")
+#Using cached local data for templates
+t3 = ToExcel(domain='mobile', server=False, local='path/to/local/stix')
+t3.to_xlsx(layer=lay, filepath="demo3.xlsx")
 ```

layers/core/gradient.py

@@ -58,17 +58,20 @@ def _compute_curve(self):
             Computes the gradient color curve
         """
         if self.maxValue is not None and self.minValue is not None and self.colors is not None:
-            if len(self.colors) == 2:
-                s_c = colour.Color(self.colors[0])
-                e_c = colour.Color(self.colors[1])
-                self.curve = list(s_c.range_to(e_c, self.maxValue - self.minValue))
-            else:
-                s_c = colour.Color(self.colors[0])
-                m_c = colour.Color(self.colors[1])
-                e_c = colour.Color(self.colors[2])
-                self.curve = list(s_c.range_to(m_c, int(math.floor(self.maxValue - self.minValue)/2)))
-                curve_2 = list(m_c.range_to(e_c, int(math.ceil(self.maxValue - self.minValue)/2)))
+            chunksize = int(math.floor((self.maxValue - self.minValue)/(len(self.colors) - 1)))
+            fchunksize = int(math.ceil((self.maxValue - self.minValue)/(len(self.colors) - 1)))
+            self.curve = []
+            index = 1
+            while index < len(self.colors):
+                s_c = colour.Color(self.colors[index-1])
+                e_c = colour.Color(self.colors[index])
+                if index == len(self.colors):
+                    curve_2 = list(s_c.range_to(e_c, fchunksize))
+                else:
+                    curve_2 = list(s_c.range_to(e_c, chunksize))
+                index += 1
                 self.curve.extend(curve_2)
+            self.curve.append(colour.Color(self.colors[-1]))
 
     def compute_color(self, score):
         """
@@ -77,9 +80,9 @@ def compute_color(self, score):
                 the gradient
         """
         if score <= self.minValue:
-            return self.colors[0]
+            return self.curve[0].hex_l
         if score >= self.maxValue:
-            return self.colors[-1]
+            return self.curve[-1].hex_l
 
         target = self.curve[score - self.minValue]
         return target.hex_l

layers/exporters/excel_templates.py

@@ -11,33 +11,37 @@ class BadTemplateException(Exception):
 
 class ExcelTemplates:
 
-    def __init__(self, fresh=False, mode='enterprise'):
+    def __init__(self, server=False, local=None, domain='enterprise'):
         """
             Initialization - Creates a ExcelTemplate object
 
-            :param mode: The domain to utilize
+            :param server: Whether or not to use the taxii server to build the matrix
+            :param local: Optional path to local taxii data
+            :param domain: The domain to utilize
         """
-        muse = mode
+        muse = domain
         if muse.startswith('mitre-'):
-            muse = mode[6:]
+            muse = domain[6:]
         if muse in ['enterprise', 'mobile']:
             self.mode = muse
-            h = MatrixGen(fresh=fresh)
-            self.codex = h.build_matrix(muse)
+            self.h = MatrixGen(server=server, local=local)
+            self.codex = self.h.get_matrix(muse)
         else:
             raise BadTemplateException
 
-    def _build_raw(self, showName=True, showID=False, subtechs=[], exclude=[]):
+    def _build_raw(self, showName=True, showID=False, sort=0, scores=[], subtechs=[], exclude=[]):
         """
             INTERNAL - builds a raw, not-yet-marked-up excel document based on the specifications
 
             :param showName: Whether or not to display names for each entry
             :param showID: Whether or not to display Technique IDs for each entry
+            :param sort: The sort mode to use
             :param subtechs: List of all visible subtechniques
             :param exclude: List of of techniques to exclude from the matrix
             :return: a openpyxl workbook object containing the raw matrix
         """
-        template, joins = MatrixGen._construct_panop(self.codex, subtechs, exclude)
+        self.codex = self.h._adjust_ordering(self.codex, sort, scores)
+        template, joins = self.h._construct_panop(self.codex, subtechs, exclude)
         self.template = template
         wb = openpyxl.Workbook()
 
@@ -52,11 +56,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 = MatrixGen._get_ID(self.codex, template[entry]) + ': ' + template[entry]
+                write_val = self.h._get_ID(self.codex, template[entry]) + ': ' + template[entry]
             elif showName:
                 write_val = template[entry]
             elif showID:
-                write_val = MatrixGen._get_ID(self.codex, template[entry])
+                write_val = self.h._get_ID(self.codex, template[entry])
             c.value = write_val
             if entry[0] == 1:
                 c.font = header_template_f
@@ -83,18 +87,19 @@ def _build_raw(self, showName=True, showID=False, subtechs=[], exclude=[]):
 
         return wb
 
-    def export(self, showName, showID, subtechs=[], exclude=[]):
+    def export(self, showName, showID, sort=0, scores=[], subtechs=[], exclude=[]):
         """
             Export a raw customized excel template
 
 
             :param showName: Whether or not to display names for each entry
             :param showID: Whether or not to display Technique IDs for each entry
+            :param sort: The sort mode to utilize
             :param subtechs: List of all visible subtechniques
             :param exclude: List of of techniques to exclude from the matrix
             return: a openpyxl workbook object containing the raw matrix
         """
-        return self._build_raw(showName, showID, subtechs, exclude)
+        return self._build_raw(showName, showID, sort, scores, subtechs, exclude)
 
     def retrieve_coords(self, techniqueID, tactic=None):
         """
@@ -105,16 +110,21 @@ def retrieve_coords(self, techniqueID, tactic=None):
             :return: A tuple representing the (row, column) of the target element in the workbook
         """
         listing = []
-        match = MatrixGen._get_name(self.codex, techniqueID)
+        match = self.h._get_name(self.codex, techniqueID)
         for entry in self.template:
             if self.template[entry] == match:
                 if tactic is not None:
                     try:
-                        if self.template[(1, entry[1])] != MatrixGen.convert[tactic]:
+                        if self.template[(1, entry[1])] != self.h.convert(tactic):
                             continue
                     except KeyError:
                         # account for subtechniques when scanning
-                        if self.template[(1, entry[1] - 1)] != MatrixGen.convert[tactic]:
+                        if self.template[(1, entry[1] - 1)] != self.h.convert(tactic):
                             continue
                 listing.append(entry)
+        if listing == []:
+            if '.' in techniqueID:
+                parent = self.retrieve_coords(techniqueID.split('.')[0], tactic)
+                if parent != []:
+                    return 'HIDDEN'
         return listing