Adding support for multi threat libraries following issue #261

Author: izar tarandach

.gitignore

@@ -126,3 +126,4 @@ plantuml.jar
 tm/
 /sqldump
 /tests/.config.pytm
+.aider*

pytm/pytm.py

@@ -51,6 +51,7 @@ def __init__(self, e, context):
 
 
 logger = logging.getLogger(__name__)
+_defaultThreatsFile = os.path.dirname(__file__) + "/threatlib/threats.json"
 
 
 class var(object):
@@ -787,8 +788,7 @@ class TM:
     )
     name = varString("", required=True, doc="Model name")
     description = varString("", required=True, doc="Model description")
-    threatsFile = varStrings("")
-    threatsFileInit = False
+    threatsFile = varStrings([_defaultThreatsFile])
     isOrdered = varBool(False, doc="Automatically order all Dataflows")
     mergeResponses = varBool(False, doc="Merge response edges in DFDs")
     ignoreUnused = varBool(False, doc="Ignore elements not used in any Dataflow")
@@ -836,9 +836,9 @@ def _init_threats(self):
 
     def _add_threats(self):
 
-        for threat_file in self.threatsFile:
+        for tf in self.threatsFile:
             try:
-                with open(threat_file, "r", encoding="utf8") as threat_file:
+                with open(tf, "r", encoding="utf8") as threat_file:
                     threats_json = json.load(threat_file)
             except (FileNotFoundError, PermissionError, IsADirectoryError) as e:
                 raise UIError(
@@ -1131,14 +1131,19 @@ def _process(self):
 
         # delaying loading of threats to accomodate multiple threat files in the
         # command line
-        if self.threatsFileInit == False:
-            tfs = [os.path.dirname(__file__) + "/threatlib/threats.json"]
-            if result.threat_files:
-                tfs.extend(result.threat_files)
-            self.threatsFile = tfs
-            self._init_threats()
-            self.threatsFileInit = True
-
+        if result.threat_files:
+            # start by removing the default
+            del self.threatsFile[0]
+            if "default" in result.threat_files:
+                index = result.threat_files.index("default")
+                result.threat_files[index] = _defaultThreatsFile
+            for x in result.threat_files:
+                self.threatsFile.append(x)
+        else:
+            # it is just the default file, so no need to do anything
+            pass
+        self._init_threats()
+            
         if result.debug:
             logger.setLevel(logging.DEBUG)
 

sample.png

@@ -1405,5 +1405,5 @@
     "name": "my test tm",
     "onDuplicates": "Action.NO_ACTION",
     "threatsExcluded": [],
-    "threatsFile": "pytm/threatlib/threats.json"
+    "threatsFile": "{'pytm/threatlib/threats.json'}"
 }

tests/output.json

@@ -49,12 +49,16 @@ def test_kwargs(self):
 
     def test_load_threats(self):
         tm = TM("TM")
-        self.assertNotEqual(len(TM._threats), 1)
+        self.assertNotEqual(len(TM._threats), 0)
+
+        ''' commenting this bit off since it is now valid to change threatsFile,
+        but looking for confirmation from Jan that it is ok to do so
         with self.assertRaises(UIError):
             tm.threatsFile = "threats.json"
 
         with self.assertRaises(UIError):
             TM("TM", threatsFile="threats.json")
+        '''
 
     def test_responses(self):
         tm = TM("my test tm", description="aa", isOrdered=True)