Repository: plexsystems/konstraint
Pull request #17 (2020-07-18 20:58:52)

Refactor: Split rego functions from commands logic
By splitting the rego functions into their own package we can keep the code for the individual commands cleaner and more concise. This also lays some of the groundwork to address #8 


Patch:
diff --git a/internal/commands/create.go b/internal/commands/create.go
index b0539b4..d74b7f9 100644
--- a/internal/commands/create.go
+++ b/internal/commands/create.go
@@ -10,7 +10,7 @@ import (
 
 	"github.com/ghodss/yaml"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/apis/templates/v1beta1"
-	"github.com/open-policy-agent/opa/ast"
+	"github.com/plexsystems/konstraint/internal/rego"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -18,43 +18,6 @@ import (
 	"k8s.io/apimachinery/pkg/runtime/schema"
 )
 
-type regoFile struct {
-	filePath       string
-	packageName    string
-	importPackages []string
-	contents       string
-}
-
-func newRegoFile(filePath string, contents string) (regoFile, error) {
-	module, err := ast.ParseModule(filePath, contents)
-	if err != nil {
-		return regoFile{}, fmt.Errorf("parse module: %w", err)
-	}
-
-	var importPackages []string
-	for i := range module.Imports {
-		importPackages = append(importPackages, module.Imports[i].Path.String())
-	}
-
-	regoFile := regoFile{
-		filePath:       filePath,
-		packageName:    module.Package.Path.String(),
-		importPackages: importPackages,
-		contents:       contents,
-	}
-
-	return regoFile, nil
-}
-
-func moduleHasViolationRule(module *ast.Module) bool {
-	for _, rule := range module.Rules {
-		if strings.HasPrefix(rule.Head.String(), "violation[") {
-			return true
-		}
-	}
-	return false
-}
-
 func newCreateCommand() *cobra.Command {
 	cmd := cobra.Command{
 		Use:   "create <dir>",
@@ -102,20 +65,12 @@ func runCreateCommand(path string) error {
 		}
 	}
 
-	policyContents, err := readFilesContents(policyFilePaths)
-	if err != nil {
-		return fmt.Errorf("read policy contents: %w", err)
-	}
-	policies, err := loadPolicyFiles(policyContents)
+	policies, err := rego.LoadPoliciesWithAction(policyFilePaths, "violation")
 	if err != nil {
 		return fmt.Errorf("load policies: %w", err)
 	}
 
-	libraryContents, err := readFilesContents(libraryFilePaths)
-	if err != nil {
-		return fmt.Errorf("read library contents: %w", err)
-	}
-	libraries, err := loadLibraryFiles(libraryContents)
+	libraries, err := rego.LoadLibraries(libraryFilePaths)
 	if err != nil {
 		return fmt.Errorf("load libraries: %w", err)
 	}
@@ -130,13 +85,13 @@ func runCreateCommand(path string) error {
 	}
 
 	for _, policy := range policies {
-		policyDir := filepath.Dir(policy.filePath)
+		policyDir := filepath.Dir(policy.FilePath)
 
 		if outputFlag == "" {
 			outputDir = policyDir
 		} else {
-			templateFileName = fmt.Sprintf("template_%s.yaml", getKindFromPath(policy.filePath))
-			constraintFileName = fmt.Sprintf("constraint_%s.yaml", getKindFromPath(policy.filePath))
+			templateFileName = fmt.Sprintf("template_%s.yaml", getKindFromPath(policy.FilePath))
+			constraintFileName = fmt.Sprintf("constraint_%s.yaml", getKindFromPath(policy.FilePath))
 		}
 
 		if _, err := os.Stat(outputDir); os.IsNotExist(err) {
@@ -176,17 +131,17 @@ func runCreateCommand(path string) error {
 	return nil
 }
 
-func getConstraintTemplate(policy regoFile, libraries []regoFile) v1beta1.ConstraintTemplate {
+func getConstraintTemplate(policy rego.File, libraries []rego.File) v1beta1.ConstraintTemplate {
 	var libs []string
-	for _, importPackage := range policy.importPackages {
+	for _, importPackage := range policy.ImportPackages {
 		for _, library := range libraries {
-			if importPackage == library.packageName {
-				libs = append(libs, library.contents)
+			if importPackage == library.PackageName {
+				libs = append(libs, library.Contents)
 			}
 		}
 	}
 
-	kind := getKindFromPath(policy.filePath)
+	kind := getKindFromPath(policy.FilePath)
 
 	constraintTemplate := v1beta1.ConstraintTemplate{
 		TypeMeta: metav1.TypeMeta{
@@ -208,7 +163,7 @@ func getConstraintTemplate(policy regoFile, libraries []regoFile) v1beta1.Constr
 				{
 					Target: "admission.k8s.gatekeeper.sh",
 					Libs:   libs,
-					Rego:   policy.contents,
+					Rego:   policy.Contents,
 				},
 			},
 		},
@@ -217,8 +172,8 @@ func getConstraintTemplate(policy regoFile, libraries []regoFile) v1beta1.Constr
 	return constraintTemplate
 }
 
-func getConstraint(policy regoFile) (unstructured.Unstructured, error) {
-	kind := getKindFromPath(policy.filePath)
+func getConstraint(policy rego.File) (unstructured.Unstructured, error) {
+	kind := getKindFromPath(policy.FilePath)
 	constraint := unstructured.Unstructured{}
 	constraint.SetName(strings.ToLower(kind))
 	constraint.SetGroupVersionKind(schema.GroupVersionKind{Group: "constraints.gatekeeper.sh", Version: "v1beta1", Kind: kind})
@@ -230,7 +185,7 @@ func getConstraint(policy regoFile) (unstructured.Unstructured, error) {
 		}
 	}
 
-	policyCommentBlocks, err := getPolicyCommentBlocks(policy.contents)
+	policyCommentBlocks, err := getPolicyCommentBlocks(policy.Contents)
 	if err != nil {
 		return unstructured.Unstructured{}, fmt.Errorf("get policy comment blocks: %w", err)
 	}
@@ -306,57 +261,6 @@ func getRegoFilePaths(path string) ([]string, error) {
 	return regoFilePaths, nil
 }
 
-func loadPolicyFiles(policyContents map[string]string) ([]regoFile, error) {
-	var policyFiles []regoFile
-	for path, contents := range policyContents {
-		module, err := ast.ParseModule(path, contents)
-		if err != nil {
-			return nil, fmt.Errorf("parse module: %w", err)
-		}
-
-		if !moduleHasViolationRule(module) {
-			continue
-		}
-
-		policyFile, err := newRegoFile(path, contents)
-		if err != nil {
-			return nil, fmt.Errorf("new rego file: %w", err)
-		}
-
-		policyFiles = append(policyFiles, policyFile)
-	}
-
-	return policyFiles, nil
-}
-
-func loadLibraryFiles(libraryContents map[string]string) ([]regoFile, error) {
-	var libraryFiles []regoFile
-	for path, contents := range libraryContents {
-		libraryFile, err := newRegoFile(path, contents)
-		if err != nil {
-			return nil, fmt.Errorf("new rego file: %w", err)
-		}
-
-		libraryFiles = append(libraryFiles, libraryFile)
-	}
-
-	return libraryFiles, nil
-}
-
-func readFilesContents(filePaths []string) (map[string]string, error) {
-	filesContents := make(map[string]string)
-	for _, filePath := range filePaths {
-		data, err := ioutil.ReadFile(filePath)
-		if err != nil {
-			return nil, fmt.Errorf("read file: %w", err)
-		}
-
-		filesContents[filePath] = string(data)
-	}
-
-	return filesContents, nil
-}
-
 func getKindFromPath(path string) string {
 	kind := filepath.Base(filepath.Dir(path))
 	kind = strings.ReplaceAll(kind, "-", " ")
diff --git a/internal/rego/rego.go b/internal/rego/rego.go
new file mode 100644
index 0000000..70579aa
--- /dev/null
+++ b/internal/rego/rego.go
@@ -0,0 +1,154 @@
+package rego
+
+import (
+	"fmt"
+	"io/ioutil"
+	"regexp"
+
+	"github.com/open-policy-agent/opa/ast"
+)
+
+// File is a parsed Rego file
+type File struct {
+	FilePath       string
+	PackageName    string
+	ImportPackages []string
+	Contents       string
+	RulesActions   []string
+	Comments       []string
+}
+
+// LoadPoliciesWithAction loads all policies from rego with rules with a given action name
+func LoadPoliciesWithAction(files []string, action string) ([]File, error) {
+	regoFiles, err := loadRegoFiles(files)
+	if err != nil {
+		return nil, fmt.Errorf("load rego files: %w", err)
+	}
+
+	policies := getPoliciesWithAction(regoFiles, action)
+	return policies, nil
+}
+
+// LoadPolicies loads all policies from rego with rules
+func LoadPolicies(files []string) ([]File, error) {
+	regoFiles, err := loadRegoFiles(files)
+	if err != nil {
+		return nil, fmt.Errorf("load rego files: %w", err)
+	}
+
+	policies := getPolicies(regoFiles)
+	return policies, nil
+}
+
+// LoadLibraries loads all libraries from rego
+func LoadLibraries(files []string) ([]File, error) {
+	regoFiles, err := loadRegoFiles(files)
+	if err != nil {
+		return nil, fmt.Errorf("load rego files: %w", err)
+	}
+
+	return regoFiles, nil
+}
+
+// NewRegoFile parses the rego and creates a File
+func NewRegoFile(filePath string, contents string) (File, error) {
+	module, err := ast.ParseModule(filePath, contents)
+	if err != nil {
+		return File{}, fmt.Errorf("parse module: %w", err)
+	}
+
+	var importPackages []string
+	for i := range module.Imports {
+		importPackages = append(importPackages, module.Imports[i].Path.String())
+	}
+
+	file := File{
+		FilePath:       filePath,
+		PackageName:    module.Package.Path.String(),
+		ImportPackages: importPackages,
+		Contents:       contents,
+		RulesActions:   getModuleRulesActions(module),
+		Comments:       getModuleComments(module),
+	}
+
+	return file, nil
+}
+
+func getPoliciesWithAction(regoFiles []File, action string) []File {
+	var matchingPolicies []File
+	allPolicies := getPolicies(regoFiles)
+	for _, policy := range allPolicies {
+		for _, ruleAction := range policy.RulesActions {
+			if ruleAction == action {
+				matchingPolicies = append(matchingPolicies, policy)
+			}
+		}
+	}
+
+	return matchingPolicies
+}
+
+func getPolicies(regoFiles []File) []File {
+	var policies []File
+	for _, regoFile := range regoFiles {
+		if len(regoFile.RulesActions) > 0 {
+			policies = append(policies, regoFile)
+		}
+	}
+
+	return policies
+}
+
+func loadRegoFiles(files []string) ([]File, error) {
+	filesContents, err := readFilesContents(files)
+	if err != nil {
+		return nil, fmt.Errorf("read files: %w", err)
+	}
+
+	var regoFiles []File
+	for path, contents := range filesContents {
+		regoFile, err := NewRegoFile(path, contents)
+		if err != nil {
+			return nil, fmt.Errorf("new rego file: %w", err)
+		}
+
+		regoFiles = append(regoFiles, regoFile)
+	}
+
+	return regoFiles, nil
+}
+
+func getModuleRulesActions(module *ast.Module) []string {
+	var rulesActions []string
+	re := regexp.MustCompile("^\\s*([a-z]+)\\s*\\[\\s*msg")
+	for _, rule := range module.Rules {
+		match := re.FindStringSubmatch(rule.Head.String())
+		if len(match) == 0 {
+			continue
+		}
+		rulesActions = append(rulesActions, match[1])
+	}
+	return rulesActions
+}
+
+func getModuleComments(module *ast.Module) []string {
+	var comments []string
+	for _, comment := range module.Comments {
+		comments = append(comments, string(comment.Text))
+	}
+	return comments
+}
+
+func readFilesContents(filePaths []string) (map[string]string, error) {
+	filesContents := make(map[string]string)
+	for _, filePath := range filePaths {
+		data, err := ioutil.ReadFile(filePath)
+		if err != nil {
+			return nil, fmt.Errorf("read file: %w", err)
+		}
+
+		filesContents[filePath] = string(data)
+	}
+
+	return filesContents, nil
+}


Test patch:
diff --git a/internal/commands/create_test.go b/internal/commands/create_test.go
index c9954dd..d4dc510 100644
--- a/internal/commands/create_test.go
+++ b/internal/commands/create_test.go
@@ -4,19 +4,21 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/plexsystems/konstraint/internal/rego"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
 
 func TestGetConstraint_NoKinds_ReturnsEmptyMatcher(t *testing.T) {
 	policy := `package test
-# Description`
+# Description
+rule[msg] { msg = true }`
 
-	rego, err := newRegoFile("/foo/test-kind", policy)
+	parsedPolicy, err := rego.NewRegoFile("test.rego", policy)
 	if err != nil {
 		t.Fatal("new rego file:", err)
 	}
 
-	actual, err := getConstraint(rego)
+	actual, err := getConstraint(parsedPolicy)
 	if err != nil {
 		t.Fatal("get constraint:", err)
 	}
@@ -29,11 +31,12 @@ func TestGetConstraint_NoKinds_ReturnsEmptyMatcher(t *testing.T) {
 func TestGetConstraint_KindsInComment_ReturnsKinds(t *testing.T) {
 	policy := `package test
 # Description
-# @Kinds core/Pod apps/Deployment`
+# @Kinds core/Pod apps/Deployment
+rule[msg] { msg = true }`
 
-	rego, err := newRegoFile("/foo/test-kind", policy)
+	rego, err := rego.NewRegoFile("test.rego", policy)
 	if err != nil {
-		t.Fatal("new rego file:", err)
+		t.Fatal("load policy rego file:", err)
 	}
 
 	actual, err := getConstraint(rego)
@@ -68,20 +71,21 @@ func TestGetConstraint_KindsInComment_ReturnsKinds(t *testing.T) {
 func TestGetConstraintTemplate_CorrectLibrariesImported(t *testing.T) {
 	policyImportsFoo := `package test
 
-import data.lib.foo`
+import data.lib.foo
+rule[msg] { msg = true }`
 
 	libraryRegos := []string{`package lib.foo`, `package lib.bar`}
 
-	policyFile, err := newRegoFile("/foo/test-kind/src.rego", policyImportsFoo)
+	policyFile, err := rego.NewRegoFile("/foo/test-kind/src.rego", policyImportsFoo)
 	if err != nil {
 		t.Fatal("new rego file:", err)
 	}
 
-	var libraries []regoFile
+	var libraries []rego.File
 	for key, library := range libraryRegos {
 		libraryPath := fmt.Sprintf("/foo/lib/library-%v.rego", key)
 
-		libraryFile, err := newRegoFile(libraryPath, library)
+		libraryFile, err := rego.NewRegoFile(libraryPath, library)
 		if err != nil {
 			t.Fatal("new rego file:", err)
 		}
@@ -104,44 +108,6 @@ import data.lib.foo`
 	}
 }
 
-func TestLoadPolicyFiles(t *testing.T) {
-	policyContents := make(map[string]string)
-	policyContents["missingViolation.rego"] = `package test
-default a = true`
-	policyContents["withViolation.rego"] = `package test
-violation[msg] {
-	msg = "test"
-}`
-
-	policies, err := loadPolicyFiles(policyContents)
-	if err != nil {
-		t.Fatal("load policy files:", err)
-	}
-
-	if len(policies) != 1 {
-		t.Error("incorrect number of policies loaded")
-	}
-}
-
-func TestLoadLibraryFiles(t *testing.T) {
-	libraryContents := make(map[string]string)
-	libraryContents["missingViolation.rego"] = `package test
-default a = true`
-	libraryContents["withViolation.rego"] = `package test
-violation[msg] {
-	msg = "test"
-}`
-
-	policies, err := loadLibraryFiles(libraryContents)
-	if err != nil {
-		t.Fatal("load library files:", err)
-	}
-
-	if len(policies) != 2 {
-		t.Error("incorrect number of libraries loaded")
-	}
-}
-
 func TestGetKindFromPath(t *testing.T) {
 	path := "/path/to/rego/container-resource-limits/something.rego"
 
diff --git a/internal/rego/rego_test.go b/internal/rego/rego_test.go
new file mode 100644
index 0000000..dafd10c
--- /dev/null
+++ b/internal/rego/rego_test.go
@@ -0,0 +1,68 @@
+package rego
+
+import (
+	"reflect"
+	"testing"
+)
+
+var (
+	testRegoFiles = []File{
+		{
+			FilePath:     "missingViolation.rego",
+			Contents:     "package test\ndefault a = true",
+			RulesActions: nil,
+		},
+		{
+			FilePath:     "withViolation.rego",
+			Contents:     "package test\nviolation[msg] { msg = true }",
+			RulesActions: []string{"violation"},
+		},
+		{
+			FilePath:     "withWarn.rego",
+			Contents:     "package test\nwarn[msg] { msg = true }",
+			RulesActions: []string{"warn"},
+		},
+	}
+)
+
+func TestGetModulesRulesActions(t *testing.T) {
+	var rulesActionsTests = []struct {
+		policy      string
+		ruleCount   int
+		ruleActions []string
+	}{
+		{"package test\ndefault test = true", 0, nil},
+		{"package test\nviolation[msg] { msg = true }", 1, []string{"violation"}},
+		{"package test\nwarn[msg] { msg = true }", 1, []string{"warn"}},
+		{"package test\nviolation[msg] { msg = true }\nwarn[msg] { msg = true }", 2, []string{"violation", "warn"}},
+	}
+
+	for _, test := range rulesActionsTests {
+		regoFile, err := NewRegoFile("test.rego", test.policy)
+		if err != nil {
+			t.Fatal("newRegoFile")
+		}
+
+		if len(regoFile.RulesActions) != test.ruleCount {
+			t.Error("incorrect rule count")
+		}
+
+		if !reflect.DeepEqual(regoFile.RulesActions, test.ruleActions) {
+			t.Errorf("incorrect rule actions\nreceived: %v\nexpected: %v", regoFile.RulesActions, test.ruleActions)
+		}
+	}
+}
+
+func TestGetPolicies(t *testing.T) {
+	policies := getPolicies(testRegoFiles)
+	if len(policies) != 2 {
+		t.Error("incorrect number of policies loaded")
+	}
+}
+
+func TestGetPoliciesWithAction(t *testing.T) {
+	matchingPolicies := getPoliciesWithAction(testRegoFiles, "warn")
+	if len(matchingPolicies) != 1 {
+		t.Error("incorrect number of policies loaded")
+	}
+}