Merge pull request #22 from nikpivkin/rego-s

refactor: remove rego scanner

Makefile

@@ -13,7 +13,7 @@ rego: fmt-rego test-rego
 
 .PHONY: fmt-rego
 fmt-rego:
-	opa fmt -w rules/cloud/policies
+	opa fmt -w rules/
 
 .PHONY: test-rego
 test-rego:

cmd/avd_generator/main.go

@@ -12,7 +12,7 @@ import (
 
 	"github.com/aquasecurity/defsec/pkg/framework"
 
-	_ "github.com/aquasecurity/trivy-policies/pkg/rego"
+	_ "github.com/aquasecurity/trivy-policies/pkg/rego/embed"
 	registered "github.com/aquasecurity/trivy-policies/pkg/rules"
 	"github.com/aquasecurity/trivy-policies/pkg/types"
 )

pkg/rego/build.go

@@ -0,0 +1,84 @@
+package rego
+
+import (
+	"io/fs"
+	"path/filepath"
+	"strings"
+
+	"github.com/aquasecurity/defsec/pkg/types"
+	"github.com/aquasecurity/trivy-policies/pkg/rego/schemas"
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/util"
+)
+
+func BuildSchemaSetFromPolicies(policies map[string]*ast.Module, paths []string, fsys fs.FS) (*ast.SchemaSet, bool, error) {
+	schemaSet := ast.NewSchemaSet()
+	schemaSet.Put(ast.MustParseRef("schema.input"), map[string]interface{}{}) // for backwards compat only
+	var customFound bool
+	for _, policy := range policies {
+		for _, annotation := range policy.Annotations {
+			for _, ss := range annotation.Schemas {
+				schemaName, err := ss.Schema.Ptr()
+				if err != nil {
+					continue
+				}
+				if schemaName != "input" {
+					if schema, ok := schemas.SchemaMap[types.Source(schemaName)]; ok {
+						customFound = true
+						schemaSet.Put(ast.MustParseRef(ss.Schema.String()), util.MustUnmarshalJSON([]byte(schema)))
+					} else {
+						b, err := findSchemaInFS(paths, fsys, schemaName)
+						if err != nil {
+							return schemaSet, true, err
+						}
+						if b != nil {
+							customFound = true
+							schemaSet.Put(ast.MustParseRef(ss.Schema.String()), util.MustUnmarshalJSON(b))
+						}
+					}
+				}
+			}
+		}
+	}
+
+	return schemaSet, customFound, nil
+}
+
+// findSchemaInFS tries to find the schema anywhere in the specified FS
+func findSchemaInFS(paths []string, srcFS fs.FS, schemaName string) ([]byte, error) {
+	var schema []byte
+	for _, path := range paths {
+		if err := fs.WalkDir(srcFS, sanitisePath(path), func(path string, info fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+			if info.IsDir() {
+				return nil
+			}
+			if !IsJSONFile(info.Name()) {
+				return nil
+			}
+			if info.Name() == schemaName+".json" {
+				schema, err = fs.ReadFile(srcFS, filepath.ToSlash(path))
+				if err != nil {
+					return err
+				}
+				return nil
+			}
+			return nil
+		}); err != nil {
+			return nil, err
+		}
+	}
+	return schema, nil
+}
+
+func IsJSONFile(name string) bool {
+	return strings.HasSuffix(name, ".json")
+}
+
+func sanitisePath(path string) string {
+	vol := filepath.VolumeName(path)
+	path = strings.TrimPrefix(path, vol)
+	return strings.TrimPrefix(strings.TrimPrefix(filepath.ToSlash(path), "./"), "/")
+}

pkg/rego/embed/embed.go

@@ -0,0 +1,125 @@
+package embed
+
+import (
+	"context"
+	"io/fs"
+	"path/filepath"
+	"strings"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/bundle"
+
+	"github.com/aquasecurity/trivy-policies/pkg/rego"
+	"github.com/aquasecurity/trivy-policies/pkg/rules"
+	rules2 "github.com/aquasecurity/trivy-policies/rules"
+)
+
+func init() {
+
+	modules, err := LoadEmbeddedPolicies()
+	if err != nil {
+		// we should panic as the policies were not embedded properly
+		panic(err)
+	}
+	loadedLibs, err := LoadEmbeddedLibraries()
+	if err != nil {
+		panic(err)
+	}
+	for name, policy := range loadedLibs {
+		modules[name] = policy
+	}
+
+	RegisterRegoRules(modules)
+}
+
+func RegisterRegoRules(modules map[string]*ast.Module) {
+	ctx := context.TODO()
+
+	schemaSet, _, _ := rego.BuildSchemaSetFromPolicies(modules, nil, nil)
+
+	compiler := ast.NewCompiler().
+		WithSchemas(schemaSet).
+		WithCapabilities(nil).
+		WithUseTypeCheckAnnotations(true)
+
+	compiler.Compile(modules)
+	if compiler.Failed() {
+		// we should panic as the embedded rego policies are syntactically incorrect...
+		panic(compiler.Errors)
+	}
+
+	retriever := rego.NewMetadataRetriever(compiler)
+	for _, module := range modules {
+		metadata, err := retriever.RetrieveMetadata(ctx, module)
+		if err != nil {
+			continue
+		}
+		if metadata.AVDID == "" {
+			continue
+		}
+		rules.Register(
+			metadata.ToRule(),
+			nil,
+		)
+	}
+}
+
+func LoadEmbeddedPolicies() (map[string]*ast.Module, error) {
+	return LoadPoliciesFromDirs(rules2.EmbeddedPolicyFileSystem, ".")
+}
+
+func LoadEmbeddedLibraries() (map[string]*ast.Module, error) {
+	return LoadPoliciesFromDirs(rules2.EmbeddedLibraryFileSystem, ".")
+}
+
+func IsRegoFile(name string) bool {
+	return strings.HasSuffix(name, bundle.RegoExt) && !strings.HasSuffix(name, "_test"+bundle.RegoExt)
+}
+
+func IsDotFile(name string) bool {
+	return strings.HasPrefix(name, ".")
+}
+
+func sanitisePath(path string) string {
+	vol := filepath.VolumeName(path)
+	path = strings.TrimPrefix(path, vol)
+	return strings.TrimPrefix(strings.TrimPrefix(filepath.ToSlash(path), "./"), "/")
+}
+
+func LoadPoliciesFromDirs(target fs.FS, paths ...string) (map[string]*ast.Module, error) {
+	modules := make(map[string]*ast.Module)
+	for _, path := range paths {
+		if err := fs.WalkDir(target, sanitisePath(path), func(path string, info fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+			if info.IsDir() {
+				return nil
+			}
+
+			if strings.HasSuffix(filepath.Dir(filepath.ToSlash(path)), "policies/advanced/optional") {
+				return fs.SkipDir
+			}
+
+			if !IsRegoFile(info.Name()) || IsDotFile(info.Name()) {
+				return nil
+			}
+			data, err := fs.ReadFile(target, filepath.ToSlash(path))
+			if err != nil {
+				return err
+			}
+			module, err := ast.ParseModuleWithOpts(path, string(data), ast.ParserOptions{
+				ProcessAnnotation: true,
+			})
+			if err != nil {
+				// s.debug.Log("Failed to load module: %s, err: %s", filepath.ToSlash(path), err.Error())
+				return err
+			}
+			modules[path] = module
+			return nil
+		}); err != nil {
+			return nil, err
+		}
+	}
+	return modules, nil
+}

pkg/rego/embed_test.go → pkg/rego/embed/embed_test.go

@@ -1,9 +1,9 @@
-package rego
+package embed
 
 import (
 	"testing"
 
-	"github.com/aquasecurity/trivy-policies/internal/rules"
+	"github.com/aquasecurity/trivy-policies/pkg/rules"
 	rules2 "github.com/aquasecurity/trivy-policies/rules"
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/stretchr/testify/assert"
@@ -12,7 +12,7 @@ import (
 
 func Test_EmbeddedLoading(t *testing.T) {
 
-	frameworkRules := rules.GetFrameworkRules()
+	frameworkRules := rules.GetRegistered()
 	var found bool
 	for _, rule := range frameworkRules {
 		if rule.GetRule().RegoPackage != "" {
@@ -102,7 +102,7 @@ deny[res]{
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			policies, err := RecurseEmbeddedModules(rules2.EmbeddedLibraryFileSystem, ".")
+			policies, err := LoadPoliciesFromDirs(rules2.EmbeddedLibraryFileSystem, ".")
 			require.NoError(t, err)
 			newRule, err := ast.ParseModuleWithOpts("/rules/newrule.rego", tc.inputPolicy, ast.ParserOptions{
 				ProcessAnnotation: true,