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

@@ -66,6 +66,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):
@@ -254,23 +267,29 @@ class Threat():
     mitigations = varString("")
     example = varString("")
     references = varString("")
-    target = ()
+    target = var([])
+    categories = varStrings([])
 
     def __init__(self, **kwargs):
-        self.id = kwargs['SID']
-        self.description = kwargs.get('description', '')
-        self.condition = kwargs.get('condition', 'True')
-        target = kwargs.get('target', 'Element')
-        if not isinstance(target, str) and isinstance(target, Iterable):
-            target = tuple(target)
-        else:
-            target = (target,)
+        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)
-        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', '')
+
+        categories = kwargs.get("categories", [])
+        if isinstance(categories, str) or not isinstance(categories, Iterable):
+            self.categories = [categories]
+        else:
+            self.categories = list(categories)
 
     def __repr__(self):
         return "<{0}.{1}({2}) at {3}>".format(
@@ -299,16 +318,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():
@@ -364,7 +417,7 @@ def resolve(self):
             for t in TM._BagOfThreats:
                 if not t.apply(e):
                     continue
-                f = Finding(e, t.description, t.details, t.severity, t.mitigations, t.example, t.id, t.references)
+                f = Finding(e, threat=t)
                 findings.append(f)
                 elements[e].append(f)
         self.findings = findings