model loader fix: local models

Author: lolololol

lib/sparkgap/attack.py

@@ -6,12 +6,21 @@
 import importlib.machinery
 import importlib.util
 
+if "SPARKGAP_MDLDIR" in os.environ.keys():
+  ENV_MDLDIR = os.environ["SPARKGAP_MDLDIR"]
+else:
+  ENV_MDLDIR = "models"
+
+CFG_KNOWNATTACKS = ["AES_LastRound_HD","AES_LastRound_HW","AES_SboxOutMasked_HW","AES_SboxOut_HD","AES_SboxOut_HW","AES_SboxOut_LSB","AES_TTableOut_HW","DES_SboxOut_HW","Keeloq","XorOut_HD","XorOut_HW"]
 # ['attacks/__init__.py', 'attacks/AES_SboxOut_HW.py', 'attacks/DES_SboxOut_HW.py', 'attacks/AES_TTableOut_HW.py', 'attacks/XorOut_HW.py']
 def usage():
+  global CFG_KNOWNATTACKS
   print("Attack model must be one of:")
-  fl = glob.glob("support/attacks/*.py") # imported from ../
+  for a in CFG_KNOWNATTACKS:
+    print(" - %s" % a)
+  fl = glob.glob("%s/*.py" % ENV_MDLDIR) # imported from ../
   for f in fl:
-    fx = f.replace("support/attacks/","")
+    fx = f.replace("%s/" % ENV_MDLDIR,"")
     fx2 = fx.replace(".py","")
     if fx2 == "__init__":
       continue
@@ -33,7 +42,10 @@ def fetchModel(modelname):
     return am
   print("Loading internal module from name '%s'" % modelname)
   try:
-    f = importlib.import_module("sparkgap.attacks.%s" % modelname)    
+    if modelname in CFG_KNOWNATTACKS:
+      f = importlib.import_module("sparkgap.attacks.%s" % (modelname))
+    else:
+      f = importlib.import_module("%s.%s" % (ENV_MDLDIR,modelname))
     return f.AttackModel()
   except Exception as ex:
     print(ex)