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

@@ -4,6 +4,7 @@
 import random
 import uuid
 from collections import defaultdict
+from collections.abc import Iterable
 from hashlib import sha224
 from os.path import dirname
 from re import match
@@ -58,6 +59,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):
         if not isinstance(value, Boundary):
@@ -164,41 +178,76 @@ class Threat():
     mitigations = varString("")
     example = varString("")
     references = varString("")
-    target = ()
+    target = varStrings([])
+    categories = varStrings([])
 
     ''' Represents a possible threat '''
     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']
 
-    def apply(self, target):
-        if type(self.target) is list:
-            if target.__class__.__name__ not in self.target:
-                return None
+        target = json_read["target"]
+        if isinstance(target, str) or not isinstance(target, Iterable):
+            self.target = [target]
         else:
-            if target.__class__.__name__ is not self.target:
-                return None
+            self.target = list(target)
+
+        categories = json_read["categories"]
+        if isinstance(categories, str) or not isinstance(categories, Iterable):
+            self.categories = [categories]
+        else:
+            self.categories = list(categories)
+
+    def apply(self, target):
+        if target.__class__.__name__ not in self.target:
+            return None
         return eval(self.condition)
 
 
 class Finding():
     ''' This class represents a Finding - the element in question and a description of the finding '''
-    def __init__(self, element, description, details, severity, mitigations, example, id, references):
-        self.target = element
+
+    def __init__(
+        self,
+        element,
+        description=None,
+        details=None,
+        severity=None,
+        mitigations=None,
+        example=None,
+        id=None,
+        references=None,
+        categories=None,
+        threat=None,
+    ):
+        self.target = element.name
         self.description = description
         self.details = details
         self.severity = severity
         self.mitigations = mitigations
         self.example = example
         self.id = id
         self.references = references
+        self._categories = categories
+        if threat is not None:
+            self.description = threat.description
+            self.details = threat.details
+            self.severity = threat.severity
+            self.mitigations = threat.mitigations
+            self.example = threat.example
+            self.id = threat.id
+            self.references = threat.references
+            self._categories = threat.categories
+
+    @property
+    def categories(self):
+        return ', '.join(self._categories)
 
 
 class TM():
@@ -244,10 +293,12 @@ def _add_threats(self):
 
     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.name, 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: