Add packaging

* clean up code
* add setup.py
* refactoring to support

.gitignore

@@ -2,6 +2,7 @@
 *.pyc
 *.feature
 .DS_Store
+*.egg-info
 __pycache__/
 venv/
 .vscode/

materialize_threats/materialize.py

@@ -0,0 +1,127 @@
+import json, argparse, os
+from .db import Node, Edge, dbgraph
+from .mx.utils import MxUtils
+from .gherkin_stride import create_gherkins_from_threats, create_feature_file_for_gherkins
+
+
+
+class ThreatMaterializer(object):
+
+    @classmethod
+    def get_flows_with_threats(cls):
+
+        SPOOFING = 'spoofing'
+        TAMPERING = 'tampering'
+        REPUDIATION = 'repudiation'
+        INFORMATION_DISCLOSURE = 'informationDisclosure'
+        DENIAL_OF_SERVICE = 'denialOfService'
+        ELEVATION_OF_PRIVILEGE = 'elevationOfPrivilege'
+
+        SOURCE = 'source'
+        SOURCE_ZONE = 'sourceZone'
+        DESTINATION = 'destination'
+        DESTINATION_ZONE =  'destinationZone'
+        PROCESS = 'process'
+
+        threats = { 
+            SPOOFING: [],
+            TAMPERING: [],
+            REPUDIATION: [],
+            INFORMATION_DISCLOSURE: [],
+            DENIAL_OF_SERVICE: [],
+            ELEVATION_OF_PRIVILEGE: []
+        }
+
+        Source = Node.alias()
+        Destination = Node.alias()
+        Process = Node.alias()
+
+        edgequery = (
+            Edge.select(
+                Source.label.alias(SOURCE), 
+                Source.zone.alias(SOURCE_ZONE),
+                Destination.label.alias(DESTINATION), 
+                Destination.zone.alias(DESTINATION_ZONE),
+                Process.label.alias(PROCESS)
+            )
+            .join(Source, on=(Source.id == Edge.source))
+            .switch(Process)
+            .join(Process, on=(Process.id == Edge.process))
+            .switch(Destination)
+            .join(Destination, on=(Destination.id == Edge.destination))
+        ) 
+
+        threats[ELEVATION_OF_PRIVILEGE] = list(
+            edgequery.where(
+                (Source.zone < Destination.zone)
+            ).dicts()
+        )
+
+        threats[SPOOFING] = list(
+        edgequery.where(
+                (Source.zone == 0) & 
+                (Destination.zone == 1)
+            ).dicts()
+        )
+
+        threats[TAMPERING] = list(
+            edgequery.where(
+                (Source.zone < Destination.zone)
+            ).dicts()
+        )
+
+        threats[REPUDIATION] = [threat for threat in threats[SPOOFING] if threat in threats[TAMPERING]]
+
+        threats[DENIAL_OF_SERVICE] = list(
+            edgequery.where(
+                (Source.zone == 0) & 
+                (Destination.zone == 1)
+            ).dicts()
+        )
+
+        threats[INFORMATION_DISCLOSURE] = list(
+            edgequery.where(
+                (Source.zone > Destination.zone)
+            ).dicts()
+        )
+
+        return threats
+
+    @classmethod
+    def output_threats(cls, threats):
+        print(
+            json.dumps(threats, indent=4)
+        )
+
+    @classmethod
+    def materialize(cls):
+
+        args = argparse.ArgumentParser(description="Enumerate STRIDE threats from a data flow diagram and create test case stubs")
+        args.add_argument(
+            "--diagram", 
+            default="samples/sample.drawio",
+            type=argparse.FileType('r'), 
+            help="The draw.io data flow diagram filename"
+        )
+        filename = args.parse_args().diagram.name
+
+        args.add_argument(
+            "--featurefile",
+            default=os.path.basename(filename) + ".feature",
+            type=argparse.FileType('w+'),
+            help="The feature filename to write"
+        )
+
+        graph = MxUtils.parse_from_xml(file=args.parse_args().diagram)
+        zones = dbgraph.get_node_trust_zones_from_graph(graph)
+        dbgraph.load_graph_into_db(graph, zones)
+
+        threats = cls.get_flows_with_threats()
+
+        gherkin_candidates = create_gherkins_from_threats(threats)
+        feature_file = create_feature_file_for_gherkins(filename, gherkin_candidates)
+
+        args.parse_args().featurefile.write(feature_file)
+
+        cls.output_threats(threats)  
+

materialize_threats/mx/models/UserObject.py

@@ -65,11 +65,17 @@ def get_object_type(self):
 
     @classmethod 
     def infer_type_from_node(cls, node):
+
+        TRUST_ZONE = 'text;html=1;strokeColor=#82b366;fillColor=#d5e8d4;align=center;verticalAlign=middle;whiteSpace=wrap;overflow=hidden;'
+        ELEMENT = 'rounded=0;whiteSpace=wrap;html=1;'
+        PROCESS = 'ellipse;whiteSpace=wrap;html=1;aspect=fixed;'
+        DATA_STORE = 'shape=partialRectangle;whiteSpace=wrap;html=1;left=0;right=0;fillColor=none;'
+
         types = {
-            'trust zone': 'text;html=1;strokeColor=#82b366;fillColor=#d5e8d4;align=center;verticalAlign=middle;whiteSpace=wrap;overflow=hidden;',
-            'element': 'rounded=0;whiteSpace=wrap;html=1;',
-            'process': 'ellipse;whiteSpace=wrap;html=1;aspect=fixed;',
-            'data store': "shape=partialRectangle;whiteSpace=wrap;html=1;left=0;right=0;fillColor=none;"
+            'trust zone': TRUST_ZONE,
+            'element': ELEMENT,
+            'process': PROCESS,
+            'data store': DATA_STORE
 
         }
 
@@ -78,14 +84,4 @@ def infer_type_from_node(cls, node):
                 if text.text == value: 
                     return key
 
-        return None
-
-"""
-class ZoneObject(UserObject):
-    def __init__(self, xml, label):
-        super(ZoneObject, self).__init__(xml, label)
-    
-    def get_trust_zone(self):
-        zone = self.xml.get(self.LABEL)
-        return zone.lower().split(self.ZONE_PREFIX)[self.ZONE_INDEX]
-"""
+        return None

readme.md

@@ -74,8 +74,8 @@ These are just a few ideas.
 ```
 git clone [email protected]:secmerc/materialize_threats.git
 cd materialize_threats
-pip3 install -r requirements.txt
-python3 materialize.py --filename=/path/to/diagram.drawio
+pip install -e .
+materialize-threats --diagram=/path/to/diagram.drawio
 ```
 
 ## 3. Creating the feature file
@@ -96,4 +96,7 @@ python3 materialize_threats/materialize.py --filename=samples/bookface.drawio
 # :warning: Is this production ready?
 Not yet.
 * There are no tests written, but im pretty sure it works. 
-* Lots of other python stuff that might horrify you but wont impact functionality that I know of.
+* Lots of other python stuff that might horrify you but wont impact functionality that I know of.
+
+# Handy links
+* https://docs.microsoft.com/en-us/archive/blogs/larryosterman/threat-modeling-again-presenting-the-playsound-threat-model

setup.py

@@ -0,0 +1,25 @@
+import setuptools
+
+with open("README.md", "r") as fh:
+    long_description = fh.read()
+
+setuptools.setup(
+    name="materialize-threats-secmerc", # Replace with your own username
+    version="1.0.0",
+    author="Jacob Salassi",
+    author_email="[email protected]",
+    description="Analyze draw.io data flow diagrams for STRIDE threat classes",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/secmerc/materialize_threats",
+    packages=setuptools.find_packages(),
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires='>=3.6',
+    entry_points = {
+        'console_scripts': ['materialize-threats=materialize_threats.materialize:ThreatMaterializer.materialize'],
+    }
+)