add categories to threats

docs/template.md

@@ -34,6 +34,8 @@ Name|From|To |Data|Protocol|Port
   <p> {{item.target}} </p>
   <h6> Severity </h6>
   <p>{{item.severity}}</p>
+  <h6> Categories </h6>
+  <p>{{item.categories}}</p>
   <h6>Example Instances</h6>
   <p>{{item.example}}</p>
   <h6>Mitigations</h6>

pytm/pytm.py

@@ -65,6 +65,19 @@ def __set__(self, instance, value):
         super().__set__(instance, value)
 
 
+class varStrings(var):
+
+    def __set__(self, instance, value):
+        for i, e in enumerate(value):
+            if not isinstance(e, str):
+                raise ValueError(
+                    "expecting a list of Strings, item number {} is a {}".format(
+                        i, type(value)
+                    )
+                )
+        super().__set__(instance, list(value))
+
+
 class varBoundary(var):
 
     def __set__(self, instance, value):
@@ -204,24 +217,29 @@ class Threat():
     mitigations = varString("")
     example = varString("")
     references = varString("")
-    target = ()
-
-    def __init__(self, json_read):
-        self.id = json_read['SID']
-        self.description = json_read['description']
-        self.condition = json_read['condition']
-        self.target = json_read['target']
-        self.details = json_read['details']
-        self.severity = json_read['severity']
-        self.mitigations = json_read['mitigations']
-        self.example = json_read['example']
-        self.references = json_read['references']
-
-        if not isinstance(self.target, str) and isinstance(self.target, Iterable):
-            self.target = tuple(self.target)
+    target = var([])
+    categories = varStrings([])
+
+    def __init__(self, **kwargs):
+        self.id = kwargs['SID']
+        self.description = kwargs.get('description', '')
+        self.condition = kwargs.get('condition', 'True')
+        self.details = kwargs.get('details', '')
+        self.severity = kwargs.get('severity', '')
+        self.mitigations = kwargs.get('mitigations', '')
+        self.example = kwargs.get('example', '')
+        self.references = kwargs.get('references', '')
+
+        target = kwargs.get("target", "Element")
+        if isinstance(target, str) or not isinstance(target, Iterable):
+            target = [target]
+        self.target = tuple(getattr(sys.modules[__name__], x) for x in target)
+
+        categories = kwargs.get("categories", [])
+        if isinstance(categories, str) or not isinstance(categories, Iterable):
+            self.categories = [categories]
         else:
-            self.target = (self.target,)
-        self.target = tuple(getattr(sys.modules[__name__], x) for x in self.target)
+            self.categories = list(categories)
 
     def __repr__(self):
         return "<{0}.{1}({2}) at {3}>".format(
@@ -250,16 +268,50 @@ class Finding():
     id = varString("", required=True, doc="Threat ID")
     references = varString("", required=True, doc="Threat references")
 
-    def __init__(self, element, description, details, severity, mitigations, example, id, references):
+    def __init__(
+        self,
+        element,
+        **kwargs,
+    ):
         self.target = element.name
         self.element = element
-        self.description = description
-        self.details = details
-        self.severity = severity
-        self.mitigations = mitigations
-        self.example = example
-        self.id = id
-        self.references = references
+        attrs = [
+            "description",
+            "details",
+            "severity",
+            "mitigations",
+            "example",
+            "id",
+            "references",
+        ]
+        threat = kwargs.get("threat", None)
+        if threat:
+            for a in attrs:
+                setattr(self, a, getattr(threat, a))
+            setattr(self, "_categories", threat.categories)
+            return
+
+        for a in attrs:
+            if a in kwargs:
+                setattr(self, a, kwargs.get(a))
+            if "categories" in kwargs:
+                setattr(
+                    self,
+                    "_categories",
+                    kwargs.get("categories"),
+                )
+
+    def __repr__(self):
+        return "<{0}.{1}({2}) at {3}>".format(
+            self.__module__, type(self).__name__, self.id, hex(id(self))
+        )
+
+    def __str__(self):
+        return "{0}({1})".format(type(self).__name__, self.id)
+
+    @property
+    def categories(self):
+        return ', '.join(self._categories)
 
 
 class TM():
@@ -304,14 +356,16 @@ def _add_threats(self):
             threats_json = json.load(threat_file)
 
         for i in threats_json:
-            TM._BagOfThreats.append(Threat(i))
+            TM._BagOfThreats.append(Threat(**i))
 
     def resolve(self):
         for e in (TM._BagOfElements):
-            if e.inScope is True:
-                for t in TM._BagOfThreats:
-                    if t.apply(e) is True:
-                        TM._BagOfFindings.append(Finding(e, t.description, t.details, t.severity, t.mitigations, t.example, t.id, t.references))
+            if not e.inScope:
+                continue
+            for t in TM._BagOfThreats:
+                if not t.apply(e):
+                    continue
+                TM._BagOfFindings.append(Finding(e, threat=t))
 
     def check(self):
         if self.description is None: