From 6229fd9913e1c8e0e215ccfd59b158679aee3b21 Mon Sep 17 00:00:00 2001
From: hectorj2f <hectorf@vmware.com>
Date: Fri, 22 Apr 2022 14:11:25 +0300
Subject: [PATCH] chore: add rego function to consume modules

Signed-off-by: hectorj2f <hectorf@vmware.com>
---
 pkg/cosign/rego/rego.go      |  46 ++++++++++++++
 pkg/cosign/rego/rego_test.go | 115 +++++++++++++++++++++++++++++++++++
 2 files changed, 161 insertions(+)

diff --git a/pkg/cosign/rego/rego.go b/pkg/cosign/rego/rego.go
index f3def4dcc76..2f4a6e0513a 100644
--- a/pkg/cosign/rego/rego.go
+++ b/pkg/cosign/rego/rego.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 
 	"github.com/open-policy-agent/opa/rego"
+	"knative.dev/pkg/logging"
 )
 
 // The query below should meet the following requirements:
@@ -29,6 +30,12 @@ import (
 // * Queries for a single value.
 const QUERY = "data.signature.allow"
 
+// CosignRegoPackageName defines the expected package name of a provided rego module
+const CosignRegoPackageName = "sigstore"
+
+// CosignEvaluationRule defines the expected evaluation role of a provided rego module
+const CosignEvaluationRule = "isCompliant"
+
 func ValidateJSON(jsonBody []byte, entrypoints []string) []error {
 	ctx := context.Background()
 
@@ -73,3 +80,42 @@ func ValidateJSON(jsonBody []byte, entrypoints []string) []error {
 	}
 	return errs
 }
+
+// ValidateJSONWithModuleInput takes the body of the results to evaluate and the defined module
+// in a policy to validate against the input data
+func ValidateJSONWithModuleInput(jsonBody []byte, moduleInput string) error {
+	ctx := context.Background()
+	query := fmt.Sprintf("%s = data.%s.%s", CosignEvaluationRule, CosignRegoPackageName, CosignEvaluationRule)
+	module := fmt.Sprintf("%s.rego", CosignRegoPackageName)
+
+	r := rego.New(
+		rego.Query(query),
+		rego.Module(module, moduleInput))
+
+	evalQuery, err := r.PrepareForEval(ctx)
+	if err != nil {
+		return err
+	}
+
+	var input interface{}
+	dec := json.NewDecoder(bytes.NewBuffer(jsonBody))
+	dec.UseNumber()
+	if err := dec.Decode(&input); err != nil {
+		return err
+	}
+
+	rs, err := evalQuery.Eval(ctx, rego.EvalInput(input))
+	if err != nil {
+		return err
+	}
+
+	for _, result := range rs {
+		isCompliant, ok := result.Bindings[CosignEvaluationRule].(bool)
+		if ok && isCompliant {
+			logging.FromContext(ctx).Info("Validated policy is compliant")
+			return nil
+		}
+	}
+
+	return fmt.Errorf("policy is not compliant for query '%s'", query)
+}
diff --git a/pkg/cosign/rego/rego_test.go b/pkg/cosign/rego/rego_test.go
index 875249594a1..a1fb4e3541a 100644
--- a/pkg/cosign/rego/rego_test.go
+++ b/pkg/cosign/rego/rego_test.go
@@ -98,3 +98,118 @@ func TestValidationJSON(t *testing.T) {
 		})
 	}
 }
+
+const attestationsJSONBody = `{
+	"authorityMatches": {
+	  "keyatt": {
+		"signatures": null,
+		"attestations": {
+		  "vuln-key": [
+			{
+			  "subject": "PLACEHOLDER",
+			  "issuer": "PLACEHOLDER"
+			}
+		  ]
+		}
+	  },
+	  "keysignature": {
+		"signatures": [
+		  {
+			"subject": "PLACEHOLDER",
+			"issuer": "PLACEHOLDER"
+		  }
+		],
+		"attestations": null
+	  },
+	  "keylessatt": {
+		"signatures": null,
+		"attestations": {
+		  "custom-keyless": [
+			{
+			  "subject": "PLACEHOLDER",
+			  "issuer": "PLACEHOLDER"
+			}
+		  ]
+		}
+	  },
+	  "keylesssignature": {
+		"signatures": [
+		  {
+			"subject": "PLACEHOLDER",
+			"issuer": "PLACEHOLDER"
+		  }
+		],
+		"attestations": null
+	  }
+	}
+  }`
+
+func TestValidateJSONWithModuleInput(t *testing.T) {
+	cases := []struct {
+		name     string
+		jsonBody string
+		policy   string
+		pass     bool
+		errorMsg string
+	}{
+		{
+			name:     "passing policy attestations",
+			jsonBody: attestationsJSONBody,
+			policy: `
+				package sigstore
+				default isCompliant = false
+				isCompliant {
+					attestationsKeylessATT := input.authorityMatches.keylessatt.attestations
+					count(attestationsKeylessATT) == 1
+
+					attestationsKeyATT := input.authorityMatches.keyatt.attestations
+					count(attestationsKeyATT) == 1
+
+					keylessSignature := input.authorityMatches.keylesssignature.signatures
+					count(keylessSignature) == 1
+				
+					keySignature := input.authorityMatches.keysignature.signatures
+					count(keySignature) == 1
+				}
+			`,
+			pass: true,
+		},
+		{
+			name:     "not passing policy attestations",
+			jsonBody: attestationsJSONBody,
+			policy: `
+				package sigstore
+
+				default isCompliant = false
+
+				isCompliant {
+					attestationsKeylessATT := input.authorityMatches.keylessatt.attestations
+					count(attestationsKeylessATT) == 0
+
+					attestationsKeyATT := input.authorityMatches.keyatt.attestations
+					count(attestationsKeyATT) == 1
+
+					keylessSignature := input.authorityMatches.keylesssignature.signatures
+					count(keylessSignature) == 1
+				
+					keySignature := input.authorityMatches.keysignature.signatures
+					count(keySignature) == 1
+				}
+			`,
+			pass:     false,
+			errorMsg: "policy is not compliant for query 'isCompliant = data.sigstore.isCompliant'",
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := ValidateJSONWithModuleInput([]byte(tt.jsonBody), tt.policy); (err == nil) != tt.pass {
+				t.Fatalf("Unexpected result: %v", err)
+			} else if err != nil {
+				if fmt.Sprintf("%s", err) != tt.errorMsg {
+					t.Errorf("Expected error %q, got %q", tt.errorMsg, err)
+				}
+			}
+		})
+	}
+}