fix(rego): improve commands parsing

.gitignore

@@ -1,2 +1,3 @@
 .idea
 bundle.tar.gz
+opa

Makefile

@@ -9,7 +9,7 @@ rego: fmt-rego test-rego
 
 .PHONY: fmt-rego
 fmt-rego:
-	opa fmt -w checks/
+	opa fmt -w lib/ checks/
 
 .PHONY: test-rego
 test-rego:

checks/docker/update_instruction_alone.rego

@@ -45,7 +45,7 @@ package_managers = {
 deny[res] {
 	run := docker.run[_]
 	run_cmd := concat(" ", run.Value)
-	cmds := regex.split(`\s*&&\s*`, run_cmd)
+	cmds := sh.parse_commands(run_cmd)
 
 	some package_manager
 	update_indexes := has_update(cmds, package_managers[package_manager])
@@ -66,7 +66,7 @@ update_followed_by_install(cmds, package_manager, update_indexes) {
 
 contains_cmd_with_package_manager(cmds, cmds_to_check, package_manager) = cmd_indexes {
 	cmd_indexes = [idx |
-		cmd_parts := split(cmds[idx], " ")
+		cmd_parts := cmds[idx]
 		some i, j
 		i != j
 		cmd_parts[i] == package_manager[_]

checks/docker/update_instruction_alone_test.rego

@@ -125,6 +125,25 @@ test_allowed {
 	count(r) == 0
 }
 
+test_allowed_cmds_separated_by_semicolon {
+	r := deny with input as {"Stages": [{"Name": "ubuntu:18.04", "Commands": [
+		{
+			"Cmd": "from",
+			"Value": ["ubuntu:18.04"],
+		},
+		{
+			"Cmd": "run",
+			"Value": ["apt-get update -y ; apt-get install -y curl"],
+		},
+		{
+			"Cmd": "entrypoint",
+			"Value": ["mysql"],
+		},
+	]}]}
+
+	count(r) == 0
+}
+
 test_allowed_multiple_install_cmds {
 	r := deny with input as {"Stages": [{"Name": "ubuntu:18.04", "Commands": [
 		{

cmd/opa/main.go

@@ -5,11 +5,13 @@ import (
 	"os"
 
 	// register Built-in Functions from defsec
+	"github.com/aquasecurity/trivy-checks/pkg/rego"
 	_ "github.com/aquasecurity/trivy/pkg/iac/rego"
 	"github.com/open-policy-agent/opa/cmd"
 )
 
 func main() {
+	rego.RegisterBuiltins()
 	// runs: opa test lib/ checks/
 	if err := cmd.RootCommand.Execute(); err != nil {
 		fmt.Println(err)

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/stretchr/testify v1.9.0
 	github.com/testcontainers/testcontainers-go v0.28.0
 	gopkg.in/yaml.v3 v3.0.1
+	mvdan.cc/sh/v3 v3.8.0
 )
 
 require (

go.sum

@@ -286,8 +286,8 @@ github.com/cpuguy83/dockercfg v0.3.1/go.mod h1:sugsbF4//dDlL/i+S+rtpIWp+5h0BHJHf
 github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE=
 github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
 github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
-github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY=
-github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
+github.com/creack/pty v1.1.21 h1:1/QdRyBaHHJP61QkWMXlOIBfsgdDeeKfK8SYVUWJKf0=
+github.com/creack/pty v1.1.21/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
 github.com/cyphar/filepath-securejoin v0.2.4 h1:Ugdm7cg7i6ZK6x3xDF1oEu1nfkyfH53EtKeQYTC3kyg=
 github.com/cyphar/filepath-securejoin v0.2.4/go.mod h1:aPGpWjXOXUn2NCNjFvBE6aRxGGx79pTxQpKOJNYHHl4=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -627,8 +627,8 @@ github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
-github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
-github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
+github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
+github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
 github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
@@ -1319,6 +1319,8 @@ honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
 honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
 honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
+mvdan.cc/sh/v3 v3.8.0 h1:ZxuJipLZwr/HLbASonmXtcvvC9HXY9d2lXZHnKGjFc8=
+mvdan.cc/sh/v3 v3.8.0/go.mod h1:w04623xkgBVo7/IUK89E0g8hBykgEpN0vgOj3RJr6MY=
 oras.land/oras-go/v2 v2.3.1 h1:lUC6q8RkeRReANEERLfH86iwGn55lbSWP20egdFHVec=
 oras.land/oras-go/v2 v2.3.1/go.mod h1:5AQXVEu1X/FKp1F9DMOb5ZItZBOa0y5dha0yCm4NR9c=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=

lib/sh/sh_test.rego

@@ -0,0 +1,28 @@
+package lib.sh
+
+test_parse_commands_with_ampersands {
+	cmds := sh.parse_commands("apt update && apt install curl")
+	count(cmds) == 2
+	cmds[0] == ["apt", "update"]
+	cmds[1] == ["apt", "install", "curl"]
+}
+
+test_parse_commands_empty_input {
+	cmds := sh.parse_commands("")
+	count(cmds) == 0
+}
+
+test_parse_commands_with_semicolon {
+	cmds := sh.parse_commands("apt update;apt install curl")
+	count(cmds) == 2
+	cmds[0] == ["apt", "update"]
+	cmds[1] == ["apt", "install", "curl"]
+}
+
+test_parse_commands_mixed {
+	cmds := sh.parse_commands("apt update; apt install curl && apt install git")
+	count(cmds) == 3
+	cmds[0] == ["apt", "update"]
+	cmds[1] == ["apt", "install", "curl"]
+	cmds[2] == ["apt", "install", "git"]
+}

pkg/rego/builtin.go

@@ -0,0 +1,15 @@
+package rego
+
+import (
+	"sync"
+
+	opa "github.com/open-policy-agent/opa/rego"
+)
+
+var registerOnce sync.Once
+
+func RegisterBuiltins() {
+	registerOnce.Do(func() {
+		opa.RegisterBuiltin1(shParseCommandsDecl, shParseCommandsImpl)
+	})
+}

pkg/rego/parse_commands.go

@@ -0,0 +1,73 @@
+package rego
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/topdown/builtins"
+	"github.com/open-policy-agent/opa/types"
+	"mvdan.cc/sh/v3/syntax"
+)
+
+var shParseCommandsDecl = &rego.Function{
+	Name:        "sh.parse_commands",
+	Decl:        types.NewFunction(types.Args(types.S), types.NewArray(nil, types.NewArray(nil, types.S))),
+	Description: "Parse command sequence",
+	Memoize:     true,
+}
+
+var shParseCommandsImpl = func(c rego.BuiltinContext, a *ast.Term) (*ast.Term, error) {
+	astr, err := builtins.StringOperand(a.Value, 0)
+	if err != nil {
+		return nil, fmt.Errorf("invalid parameter type: %w", err)
+	}
+
+	commands, err := parseCommands(string(astr))
+
+	if err != nil {
+		return nil, fmt.Errorf("parse command sequence error: %w", err)
+	}
+
+	var commandsTerm []*ast.Term
+	for _, cmd := range commands {
+		var cmdTerm []*ast.Term
+		for _, cmd_part := range cmd {
+			cmdTerm = append(cmdTerm, ast.StringTerm(cmd_part))
+		}
+		commandsTerm = append(commandsTerm, ast.ArrayTerm(cmdTerm...))
+	}
+
+	return ast.ArrayTerm(commandsTerm...), nil
+}
+
+func parseCommands(cmdsSeq string) ([][]string, error) {
+	f, err := syntax.NewParser().Parse(strings.NewReader(cmdsSeq), "")
+	if err != nil {
+		return nil, err
+	}
+
+	printer := syntax.NewPrinter()
+
+	var commands [][]string
+	syntax.Walk(f, func(node syntax.Node) bool {
+		switch x := node.(type) {
+		case *syntax.CallExpr:
+			args := x.Args
+			var cmd []string
+			for _, word := range args {
+				var buffer bytes.Buffer
+				printer.Print(&buffer, word)
+				cmd = append(cmd, buffer.String())
+			}
+			if cmd != nil {
+				commands = append(commands, cmd)
+			}
+		}
+		return true
+	})
+
+	return commands, nil
+}

pkg/rego/parse_commands_test.go

@@ -0,0 +1,40 @@
+package rego
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseCommands(t *testing.T) {
+	tests := []struct {
+		cmdsSeq  string
+		expected [][]string
+	}{
+		{
+			cmdsSeq:  "apt update; apt install -y nginx",
+			expected: [][]string{{"apt", "update"}, {"apt", "install", "-y", "nginx"}},
+		},
+		{
+			cmdsSeq:  "apt update && apt install -y nginx",
+			expected: [][]string{{"apt", "update"}, {"apt", "install", "-y", "nginx"}},
+		},
+		{
+			cmdsSeq:  "apt update || apt install -y nginx",
+			expected: [][]string{{"apt", "update"}, {"apt", "install", "-y", "nginx"}},
+		},
+		{
+			cmdsSeq:  `echo "test;test" ;apt update && apt install -y nginx`,
+			expected: [][]string{{"echo", "\"test;test\""}, {"apt", "update"}, {"apt", "install", "-y", "nginx"}},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.cmdsSeq, func(t *testing.T) {
+			got, err := parseCommands(test.cmdsSeq)
+			require.NoError(t, err)
+			assert.Equal(t, test.expected, got)
+		})
+	}
+}

test/docker_test.go

@@ -16,8 +16,14 @@ import (
 	"github.com/liamg/memoryfs"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	builtinrego "github.com/aquasecurity/trivy-checks/pkg/rego"
 )
 
+func init() {
+	builtinrego.RegisterBuiltins()
+}
+
 func getFileName(fpath string, info os.FileInfo, typePolicy bool) string {
 	pathParts := strings.Split(fpath, filepath.FromSlash("/"))
 	fileName := info.Name()