chore: load checks from repo (#114)

Signed-off-by: Nikita Pivkin <nikita.pivkin@smartforce.io>

Author: nikpivkin

docGen/defsec.go

@@ -11,15 +11,26 @@ import (
 	"strings"
 	"text/template"
 
+	"gopkg.in/yaml.v3"
+
 	"github.com/aquasecurity/avd-generator/menu"
 	"github.com/aquasecurity/avd-generator/util"
 	"github.com/aquasecurity/trivy/pkg/iac/framework"
 	"github.com/aquasecurity/trivy/pkg/iac/rego"
 	"github.com/aquasecurity/trivy/pkg/iac/rules"
 	"github.com/aquasecurity/trivy/pkg/iac/scan"
-	"gopkg.in/yaml.v3"
 )
 
+func registerChecks(fsys fs.FS) error {
+	rules.Reset()
+	modules, err := rego.LoadPoliciesFromDirs(fsys, "checks", "lib")
+	if err != nil {
+		return fmt.Errorf("load checks: %w", err)
+	}
+	rego.RegisterRegoRules(modules)
+	return nil
+}
+
 type DefsecComplianceSpec struct {
 	Spec struct {
 		ID               string   `yaml:"id"`
@@ -42,26 +53,12 @@ type DefsecComplianceSpec struct {
 	} `yaml:"spec"`
 }
 
-var funcMap = template.FuncMap{
-	"toLower":    strings.ToLower,
-	"toUpper":    strings.ToUpper,
-	"toTitle":    strings.Title,
-	"getSummary": getSummary,
-}
-
-var registeredRulesSummaries = make(map[string]string)
-
-func init() {
-	rules.Reset()
-
-	rego.LoadAndRegister()
+func generateDefsecComplianceSpecPages(specDir, contentDir string) {
 
+	ruleSummaries := make(map[string]string)
 	for _, rule := range rules.GetRegistered(framework.ALL) {
-		registeredRulesSummaries[rule.GetRule().AVDID] = rule.GetRule().Summary
+		ruleSummaries[rule.GetRule().AVDID] = rule.GetRule().Summary
 	}
-}
-
-func generateDefsecComplianceSpecPages(specDir, contentDir string) {
 
 	if err := filepath.Walk(specDir, func(path string, info fs.FileInfo, err error) error {
 		if err != nil {
@@ -92,27 +89,16 @@ func generateDefsecComplianceSpecPages(specDir, contentDir string) {
 			[]menu.BreadCrumb{{Name: "Compliance", Url: "/compliance"},
 				{Name: strings.Title(spec.Spec.Category), Url: fmt.Sprintf("/compliance/%s", spec.Spec.Category)}}, spec.Spec.Category, true)
 
-		return generateDefsecComplianceSpecPage(spec, contentDir)
+		return generateDefsecComplianceSpecPage(spec, contentDir, ruleSummaries)
 
 	}); err != nil {
 		fmt.Println(err)
 	}
 
 }
 
-func getSummary(id string) string {
-	if summary, ok := registeredRulesSummaries[id]; ok {
-		return fmt.Sprintf(" - %s", summary)
-	}
-
-	return ""
-
-}
-
-func generateDefsecComplianceSpecPage(spec DefsecComplianceSpec, contentDir string) error {
-
+func generateDefsecComplianceSpecPage(spec DefsecComplianceSpec, contentDir string, ruleSummaries map[string]string) error {
 	for _, control := range spec.Spec.Controls {
-
 		outputFilePath := filepath.Join(contentDir, spec.Spec.Category, fmt.Sprintf("%s-%s", spec.Spec.Title, spec.Spec.Version), fmt.Sprintf("%s.md", control.ID))
 
 		if err := os.MkdirAll(filepath.Dir(outputFilePath), 0755); err != nil {
@@ -124,7 +110,19 @@ func generateDefsecComplianceSpecPage(spec DefsecComplianceSpec, contentDir stri
 			return err
 		}
 
-		t := template.Must(template.New("defsecPost").Funcs(funcMap).Parse(defsecComplianceTemplate))
+		funcs := template.FuncMap{
+			"toLower": strings.ToLower,
+			"toUpper": strings.ToUpper,
+			"toTitle": strings.Title,
+			"getSummary": func(id string) string {
+				if summary, ok := ruleSummaries[id]; ok {
+					return fmt.Sprintf(" - %s", summary)
+				}
+				return ""
+			},
+		}
+
+		t := template.Must(template.New("defsecPost").Funcs(funcs).Parse(defsecComplianceTemplate))
 		if err := t.Execute(outputFile, map[string]interface{}{
 			"ID":          spec.Spec.ID,
 			"Version":     spec.Spec.Version,
@@ -145,7 +143,6 @@ func generateDefsecComplianceSpecPage(spec DefsecComplianceSpec, contentDir stri
 
 func generateDefsecPages(remediationDir, contentDir string) {
 	for _, r := range rules.GetRegistered(framework.ALL) {
-
 		avdId := r.GetRule().AVDID
 		topLevelID := strings.ToLower(r.GetRule().Provider.ConstName())
 		branchID := r.GetRule().Service

docGen/defsec_test.go

@@ -7,18 +7,20 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/aquasecurity/trivy/pkg/iac/rego"
 )
 
 func TestLoadsAsExpected(t *testing.T) {
+	rego.LoadAndRegister()
 
-	tempDir := t.TempDir()
-
-	generateDefsecPages("../goldens/defsec/md", tempDir)
+	outputDir := t.TempDir()
+	generateDefsecPages("../goldens/defsec/md", outputDir)
 
 	ids := []string{"avd-aws-0018"}
 
 	for _, id := range ids {
-		content, err := os.ReadFile(fmt.Sprintf("%s/aws/code-build/%s.md", tempDir, id))
+		content, err := os.ReadFile(fmt.Sprintf("%s/aws/code-build/%s.md", outputDir, id))
 		require.NoError(t, err)
 
 		expected, err := os.ReadFile(fmt.Sprintf("../goldens/defsec/expected/%s.md", id))

docGen/main.go

@@ -40,6 +40,10 @@ func main() {
 		Years = append(Years, strconv.Itoa(y))
 	}
 
+	if err := registerChecks(os.DirFS("../avd-repo/trivy-policies-repo")); err != nil {
+		fail(err)
+	}
+
 	generateChainBenchPages("../avd-repo/chain-bench-repo/internal/checks", "../avd-repo/content/compliance")
 	generateKubeBenchPages("../avd-repo/kube-bench-repo/cfg", "../avd-repo/content/compliance")
 	generateDefsecComplianceSpecPages("../avd-repo/trivy-policies-repo/pkg/specs/compliance", "../avd-repo/content/compliance")