From 9ed3368c89dfc02ee57edcec7f946a16e7e24a1b Mon Sep 17 00:00:00 2001
From: sg <spyros@ocurity.com>
Date: Thu, 10 Oct 2024 11:49:27 +0100
Subject: [PATCH] fixup-modelscan

---
 components/producers/modelscan/main.go      |  5 +++
 components/producers/modelscan/main_test.go |  8 ++---
 components/producers/modelscan/task.yaml    | 34 ++++++++++++++-------
 3 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/components/producers/modelscan/main.go b/components/producers/modelscan/main.go
index 8804d5df1..c9f69206e 100644
--- a/components/producers/modelscan/main.go
+++ b/components/producers/modelscan/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"encoding/json"
 	"log"
+	"log/slog"
 
 	v1 "github.com/ocurity/dracon/api/proto/v1"
 
@@ -38,6 +39,7 @@ func main() {
 
 func parseIssues(out *ModelScanOut) ([]*v1.Issue, error) {
 	issues := []*v1.Issue{}
+	slog.Info("found Critical issues", slog.Int("numCrit", len(out.IssuesBySeverity.Critical)))
 	for _, crit := range out.IssuesBySeverity.Critical {
 		issues = append(issues,
 			&v1.Issue{
@@ -49,6 +51,7 @@ func parseIssues(out *ModelScanOut) ([]*v1.Issue, error) {
 				Confidence:  v1.Confidence_CONFIDENCE_UNSPECIFIED,
 			})
 	}
+	slog.Info("found High issues", slog.Int("numHigh", len(out.IssuesBySeverity.High)))
 	for _, crit := range out.IssuesBySeverity.High {
 		issues = append(issues,
 			&v1.Issue{
@@ -60,6 +63,7 @@ func parseIssues(out *ModelScanOut) ([]*v1.Issue, error) {
 				Confidence:  v1.Confidence_CONFIDENCE_UNSPECIFIED,
 			})
 	}
+	slog.Info("found Medium issues", slog.Int("numMedium", len(out.IssuesBySeverity.Medium)))
 	for _, crit := range out.IssuesBySeverity.Medium {
 		issues = append(issues,
 			&v1.Issue{
@@ -71,6 +75,7 @@ func parseIssues(out *ModelScanOut) ([]*v1.Issue, error) {
 				Confidence:  v1.Confidence_CONFIDENCE_UNSPECIFIED,
 			})
 	}
+	slog.Info("found Low issues", slog.Int("numLow", len(out.IssuesBySeverity.Low)))
 	for _, crit := range out.IssuesBySeverity.Low {
 		issues = append(issues,
 			&v1.Issue{
diff --git a/components/producers/modelscan/main_test.go b/components/producers/modelscan/main_test.go
index da8c79f4f..e3920f70a 100644
--- a/components/producers/modelscan/main_test.go
+++ b/components/producers/modelscan/main_test.go
@@ -51,12 +51,12 @@ const modelScanOut = `{
   "modelscan_version": "0.5.0",
   "timestamp": "2024-01-25T17:56:00.855056",
   "input_path": "/Users/mehrinkiani/Documents/modelscan/notebooks/XGBoostModels/unsafe_model.pkl",
-  "total_issues": 1,
+  "total_issues": 4,
   "summary": {
     "total_issues_by_severity": {
-      "LOW": 0,
-      "MEDIUM": 0,
-      "HIGH": 0,
+      "LOW": 1,
+      "MEDIUM": 1,
+      "HIGH": 1,
       "CRITICAL": 1
     }
   },
diff --git a/components/producers/modelscan/task.yaml b/components/producers/modelscan/task.yaml
index 3222125a4..a3578c180 100644
--- a/components/producers/modelscan/task.yaml
+++ b/components/producers/modelscan/task.yaml
@@ -20,34 +20,46 @@ spec:
       description: The workspace containing the source-code to scan.
   steps:
   - name: run-modelscan
-    image: python:alpine
+    image: python:3.11-alpine
+    imagePullPolicy: Always
     script: |
-      pip install 'modelscan[ tensorflow, h5py ]'
-      modelscan \
-      --path "$(workspaces.output.path)/source-code/$(params.producer-modelscan-relative-path-to-model)" \
-      --reporting-format json \
-      --output-file /scratch/out.json
+      set -x
+      set +e
+
+      pip install 'modelscan'
+      modelscan --path "$(workspaces.output.path)/source-code/$(params.producer-modelscan-relative-path-to-model)" --reporting-format json --output-file /scratch/out.json
+
+      echo 'outfile:'
+      cat /scratch/out.json
 
       exitCode=$?
       if [[ $exitCode -eq 1 ]]; then
         echo "ModelScan found vulnerabilities"
         exit 0
-      else if [[ $exitCode -eq 2 ]]; then
+      elif [[ $exitCode -eq 2 ]]; then
         echo "ModelScan failed, error while scanning"
         exit $exitCode
-      else if [[ $exitCode -eq 3 ]]; then
+      elif [[ $exitCode -eq 3 ]]; then
         echo "ModelScan did not find any supported files while scanning"
         exit $exitCode
-      else if [[ $exitCode -eq 4 ]]; then
+      elif [[ $exitCode -eq 4 ]]; then
         echo "ModelScan encountered an error whle parsing CLI variables, the task definition has a bug"
         exit $exitCode
+      elif [[ $exitCode -eq 0 ]]; then
+        echo "ModelScan did not find any vulnerabilities"
+        exit $exitCode
+      else
+        echo "Received unexpected exit code, exiting"
+        exit $exitCode
       fi
     volumeMounts:
     - mountPath: /scratch
       name: scratch
   - name: produce-issues
-    imagePullPolicy: IfNotPresent
-    image: '{{ default "ghcr.io/ocurity/dracon" .Values.image.registry }}/components/producers/modelscan:{{ .Chart.AppVersion }}'
+    imagePullPolicy: Always
+    # image: '{{ default "ghcr.io/ocurity/dracon" .Values.image.registry }}/components/producers/modelscan:{{ .Chart.AppVersion }}'
+    image: kind-registry:5000/components/producers/modelscan:v0.52.1-dad26e45-dirty   
+
     command: ["/app/components/producers/modelscan/modelscan-parser"]
     args:
     - "-in=/scratch/out.json"