From 93b322a8d22c41f8d94553df54cff348419a6b85 Mon Sep 17 00:00:00 2001
From: Sam Washko <sam@protectai.com>
Date: Mon, 15 Apr 2024 14:36:32 -0700
Subject: [PATCH] adds back supported_extensions

---
 modelscan/settings.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/modelscan/settings.py b/modelscan/settings.py
index f4bfa9e..5f4e6ed 100644
--- a/modelscan/settings.py
+++ b/modelscan/settings.py
@@ -15,18 +15,22 @@
     "scanners": {
         "modelscan.scanners.H5LambdaDetectScan": {
             "enabled": True,
+            "supported_extensions": [".h5"],
         },
         "modelscan.scanners.KerasLambdaDetectScan": {
             "enabled": True,
+            "supported_extensions": [".keras"],
         },
         "modelscan.scanners.SavedModelLambdaDetectScan": {
             "enabled": True,
+            "supported_extensions": [".pb"],
             "unsafe_keras_operators": {
                 "Lambda": "MEDIUM",
             },
         },
         "modelscan.scanners.SavedModelTensorflowOpScan": {
             "enabled": True,
+            "supported_extensions": [".pb"],
             "unsafe_tf_operators": {
                 "ReadFile": "HIGH",
                 "WriteFile": "HIGH",
@@ -34,12 +38,22 @@
         },
         "modelscan.scanners.NumpyUnsafeOpScan": {
             "enabled": True,
+            "supported_extensions": [".npy"],
         },
         "modelscan.scanners.PickleUnsafeOpScan": {
             "enabled": True,
+            "supported_extensions": [
+                ".pkl",
+                ".pickle",
+                ".joblib",
+                ".dill",
+                ".dat",
+                ".data",
+            ],
         },
         "modelscan.scanners.PyTorchUnsafeOpScan": {
             "enabled": True,
+            "supported_extensions": [".bin", ".pt", ".pth", ".ckpt"],
         },
     },
     "middlewares": {