From cd476cd68bfa11ba971b84b4fe56f1c56fcb9dea Mon Sep 17 00:00:00 2001
From: Amit Schendel <amitschendel@gmail.com>
Date: Sun, 28 Jan 2024 15:06:39 +0200
Subject: [PATCH] Fixing network scanner rego

Signed-off-by: Amit Schendel <amitschendel@gmail.com>
---
 core/pkg/opaprocessor/networkscanner.go      | 11 ++---------
 core/pkg/opaprocessor/networkscanner_test.go | 14 +++++++++++---
 core/pkg/opaprocessor/utils.go               | 14 ++++++++++----
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/core/pkg/opaprocessor/networkscanner.go b/core/pkg/opaprocessor/networkscanner.go
index ec28bdf5..52531927 100644
--- a/core/pkg/opaprocessor/networkscanner.go
+++ b/core/pkg/opaprocessor/networkscanner.go
@@ -1,19 +1,12 @@
 package opaprocessor
 
 import (
-	"strconv"
-
 	servicediscovery "github.com/kubescape/kubescape-network-scanner/cmd"
 )
 
 // Check if the service is unauthenticated using kubescape-network-scanner.
-func isUnauthenticatedService(host, port string) bool {
-	portInt, err := strconv.Atoi(port)
-	if err != nil {
-		return false
-	}
-
-	discoveryResults, err := servicediscovery.ScanTargets(host, portInt)
+func isUnauthenticatedService(host string, port int) bool {
+	discoveryResults, err := servicediscovery.ScanTargets(host, port)
 	if err != nil {
 		return false
 	}
diff --git a/core/pkg/opaprocessor/networkscanner_test.go b/core/pkg/opaprocessor/networkscanner_test.go
index 166403a9..b26e3341 100644
--- a/core/pkg/opaprocessor/networkscanner_test.go
+++ b/core/pkg/opaprocessor/networkscanner_test.go
@@ -1,6 +1,7 @@
 package opaprocessor
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/alicebob/miniredis/v2"
@@ -18,11 +19,18 @@ func TestIsUnauthenticatedService(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	// Get the port as an integer
+	port, err := strconv.Atoi(s.Port())
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	// rego input
 	type args struct {
 		host string
-		port string
+		port int
 	}
+
 	tests := []struct {
 		name     string
 		args     args
@@ -33,7 +41,7 @@ func TestIsUnauthenticatedService(t *testing.T) {
 			"Unauthenticated service",
 			args{
 				host: s.Host(),
-				port: s.Port(),
+				port: port,
 			},
 			true,
 			assert.True,
@@ -42,7 +50,7 @@ func TestIsUnauthenticatedService(t *testing.T) {
 			"Authenticated service",
 			args{
 				host: s.Host(),
-				port: s.Port(),
+				port: port,
 			},
 			false,
 			assert.False,
diff --git a/core/pkg/opaprocessor/utils.go b/core/pkg/opaprocessor/utils.go
index a51d2714..eebf580b 100644
--- a/core/pkg/opaprocessor/utils.go
+++ b/core/pkg/opaprocessor/utils.go
@@ -115,18 +115,24 @@ var imageNameNormalizeDefinition = func(bctx rego.BuiltinContext, a *ast.Term) (
 
 var unauthenticatedServiceDeclaration = &rego.Function{
 	Name:    "networkscanner.isUnauthenticatedService",
-	Decl:    types.NewFunction(types.Args(types.S, types.A), types.B),
+	Decl:    types.NewFunction(types.Args(types.S, types.N), types.B),
 	Memoize: true,
 }
 
 var unauthenticatedServiceDefinition = func(bctx rego.BuiltinContext, a, b *ast.Term) (*ast.Term, error) {
-	aStr, err := builtins.StringOperand(a.Value, 1)
+	service, err := builtins.StringOperand(a.Value, 1)
 	if err != nil {
 		return nil, fmt.Errorf("invalid parameter type: %v", err)
 	}
-	bStr, err := builtins.StringOperand(b.Value, 1)
+	bNum, err := builtins.NumberOperand(b.Value, 1)
 	if err != nil {
 		return nil, fmt.Errorf("invalid parameter type: %v", err)
 	}
-	return ast.BooleanTerm(isUnauthenticatedService(string(aStr), string(bStr))), nil
+
+	portNumber, ok := bNum.Int()
+	if !ok {
+		return nil, fmt.Errorf("invalid parameter type: %v", err)
+	}
+
+	return ast.BooleanTerm(isUnauthenticatedService(string(service), portNumber)), nil
 }