ENG-13633: Initial discovery and classification implementation

.github/workflows/docker.yaml

@@ -0,0 +1,16 @@
+name: Build CLI Docker Image
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build Docker image
+        run: docker build . --file Dockerfile

.github/workflows/rego.yaml

@@ -0,0 +1,32 @@
+name: Rego Lint and Test
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  lint-and-test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@v4
+
+      - name: Setup Regal
+        uses: StyraInc/setup-regal@v1
+        with:
+          version: latest
+
+      - name: Lint
+        # Disable the line-length check for now, as it's too strict: it flags
+        # the long regexes in various classification policies.
+        run: regal lint --format=github --disable=line-length ./classification/labels/
+
+      - name: Setup OPA
+        uses: open-policy-agent/setup-opa@v2
+        with:
+          version: latest
+
+      - name: Run OPA Tests
+        run: opa test ./classification/labels/*.rego -v

.gitignore

@@ -253,3 +253,6 @@ __debug_bin
 
 # Backup files
 *~
+
+# Other
+/.run/

Dockerfile

@@ -0,0 +1,20 @@
+FROM golang:1.22 as build
+
+# Set destination for COPY.
+WORKDIR /app
+
+# Download dependencies.
+COPY go.mod go.sum ./
+RUN go mod download
+
+# Copy the source code.
+COPY . .
+
+# Build.
+RUN CGO_ENABLED=0 go build -ldflags="-X main.version=$(git describe --tags --always)" -o dmap cmd/*.go
+
+FROM gcr.io/distroless/static-debian12:nonroot
+
+COPY --from=build /app/dmap /dmap
+
+ENTRYPOINT ["/dmap"]

Makefile

@@ -15,3 +15,12 @@ integration-test:
 
 clean:
 	go clean -i ./...
+
+opt-fmt:
+	opa fmt --write ./classification/labels
+
+opa-lint:
+	regal lint --disable=line-length ./classification/labels/
+
+opa-test:
+	opa test ./classification/labels/*.rego -v

aws/scanner_test.go

@@ -159,19 +159,19 @@ func (s *AWSScannerTestSuite) TestScan() {
 		awsClientConstructor: func(awsConfig aws.Config) *awsClient {
 			return &awsClient{
 				config: awsConfig,
-				rds: &mock.MockRDSClient{
+				rds: &mock.RDSClient{
 					DBClusters:  s.dummyRDSClusters,
 					DBInstances: s.dummyRDSInstances,
 				},
-				redshift: &mock.MockRedshiftClient{
+				redshift: &mock.RedshiftClient{
 					Clusters: s.dummyRedshiftClusters,
 				},
-				dynamodb: &mock.MockDynamoDBClient{
+				dynamodb: &mock.DynamoDBClient{
 					TableNames: s.dummyDynamoDBTableNames,
 					Table:      s.dummyDynamoDBTable,
 					Tags:       s.dummyDynamoDBTags,
 				},
-				s3: &mock.MockS3Client{
+				s3: &mock.S3Client{
 					Buckets: s.dummyS3Buckets,
 					Tags:    s.dummyS3Tags,
 				},
@@ -330,23 +330,23 @@ func (s *AWSScannerTestSuite) TestScan_WithErrors() {
 		awsClientConstructor: func(awsConfig aws.Config) *awsClient {
 			return &awsClient{
 				config: awsConfig,
-				rds: &mock.MockRDSClient{
+				rds: &mock.RDSClient{
 					Errors: map[string]error{
 						"DescribeDBClusters":  dummyError,
 						"DescribeDBInstances": dummyError,
 					},
 				},
-				redshift: &mock.MockRedshiftClient{
+				redshift: &mock.RedshiftClient{
 					Errors: map[string]error{
 						"DescribeClusters": dummyError,
 					},
 				},
-				dynamodb: &mock.MockDynamoDBClient{
+				dynamodb: &mock.DynamoDBClient{
 					Errors: map[string]error{
 						"ListTables": dummyError,
 					},
 				},
-				s3: &mock.MockS3Client{
+				s3: &mock.S3Client{
 					Errors: map[string]error{
 						"ListBuckets":      dummyError,
 						"GetBucketTagging": dummyError,

classification/classification.go

@@ -0,0 +1,51 @@
+// Package classification provides various types and functions to facilitate
+// data classification. The type Classifier provides an interface which takes
+// arbitrary data as input and returns a classified version of that data as
+// output. The package contains at least one implementation which uses OPA and
+// Rego to perform the actual classification logic (see LabelClassifier),
+// however other implementations may be added in the future.
+package classification
+
+import (
+	"context"
+	"encoding/json"
+)
+
+// Classifier is an interface that represents a data classifier. A classifier
+// takes a set of data attributes and classifies them into a set of labels.
+type Classifier interface {
+	// Classify takes the given input, which amounts to essentially a "row of
+	// data", and returns the data classifications for that input. The input is
+	// a map of attribute names (i.e. columns) to their values. The returned
+	// Result is a map of attribute names to the set of labels that attributes
+	// were classified as.
+	Classify(ctx context.Context, input map[string]any) (Result, error)
+}
+
+// Result represents the classifications for a set of data attributes. The key
+// is the attribute (i.e. column) name and the value is the set of labels
+// that attribute was classified as.
+type Result map[string]LabelSet
+
+// LabelSet is a set of unique label names.
+type LabelSet map[string]struct{}
+
+// MarshalJSON marshals the LabelSet into a JSON array of strings, where each
+// string is the name of a label in the set.
+func (l LabelSet) MarshalJSON() ([]byte, error) {
+	keys := make([]string, 0, len(l))
+	for k := range l {
+		keys = append(keys, k)
+	}
+	return json.Marshal(keys)
+}
+
+// Classification represents the classification of a data repository attribute.
+type Classification struct {
+	// AttributePath is the full path of the data repository attribute
+	// (e.g. the column). Each element corresponds to a component, in increasing
+	// order of granularity (e.g. [database, schema, table, column]).
+	AttributePath []string `json:"attributePath"`
+	// Labels is the set of labels that the attribute was classified as.
+	Labels LabelSet `json:"labels"`
+}

classification/label.go

@@ -0,0 +1,161 @@
+package classification
+
+import (
+	"embed"
+	"errors"
+	"fmt"
+	"io/fs"
+	"os"
+	"path/filepath"
+
+	"github.com/open-policy-agent/opa/ast"
+	log "github.com/sirupsen/logrus"
+	"gopkg.in/yaml.v3"
+)
+
+var (
+	//go:embed labels/*.rego labels/labels.yaml
+	predefinedLabelsFs embed.FS
+)
+
+// InvalidLabelsError is an error type that represents an error when one or
+// more labels are invalid, e.g. they have invalid classification rules. The
+// error contains a slice of errors that caused the error, which can be
+// unwrapped to get the individual errors that caused the problems.
+type InvalidLabelsError struct {
+	errs []error
+}
+
+// Unwrap returns the errors that caused the InvalidLabelsError.
+func (e InvalidLabelsError) Unwrap() []error {
+	return e.errs
+}
+
+// Error returns a string representation of the InvalidLabelsError.
+func (e InvalidLabelsError) Error() string {
+	return errors.Join(e.errs...).Error()
+}
+
+// Label represents a data classification label.
+type Label struct {
+	// Name is the name of the label.
+	Name string `yaml:"name" json:"name"`
+	// Description is a brief description of the label.
+	Description string `yaml:"description" json:"description"`
+	// Tags are a list of arbitrary tags associated with the label.
+	Tags []string `yaml:"tags" json:"tags"`
+	// ClassificationRule is the compiled Rego classification rule used to
+	// classify data.
+	ClassificationRule *ast.Module `yaml:"-" json:"-"`
+}
+
+// NewLabel creates a new Label with the given name, description, classification
+// rule, and tags. The classification rule is expected to be the raw Rego code
+// that will be used to classify data. If the classification rule is invalid, an
+// error is returned.
+func NewLabel(name, description, classificationRule string, tags ...string) (Label, error) {
+	rule, err := parseRego(classificationRule)
+	if err != nil {
+		return Label{}, fmt.Errorf("error preparing classification rule for label %s: %w", name, err)
+	}
+	return Label{
+		Name:               name,
+		Description:        description,
+		Tags:               tags,
+		ClassificationRule: rule,
+	}, nil
+}
+
+// GetPredefinedLabels loads and returns the predefined embedded labels and
+// their classification rules. The labels are read from the embedded labels.yaml
+// file and the classification rules are read from the embedded Rego files. If
+// there is an error reading or unmarshalling the labels file, it is returned.
+// If there are errors reading or parsing a classification rules for labels, the
+// errors are aggregated into an InvalidLabelsError and returned, along with
+// the labels that were successfully read. Note that this should not return an
+// error in reality, as the embedded labels should always be valid. If it does,
+// it indicates a problem with the embedded labels!
+func GetPredefinedLabels() ([]Label, error) {
+	return getLabels("labels/labels.yaml", predefinedLabelsFs)
+}
+
+// GetCustomLabels loads and returns the labels and their classification rules
+// defined in the given labels yaml file. The labels are read from the file
+// along with their classification rule Rego files (defined in the yaml). If
+// there is an error unmarshalling the labels file, it is returned. If there are
+// errors reading or parsing a classification rules for labels, the errors are
+// aggregated into an InvalidLabelsError and returned, along with the labels
+// that were successfully read.
+func GetCustomLabels(labelsYamlFname string) ([]Label, error) {
+	path, err := filepath.Abs(labelsYamlFname)
+	if err != nil {
+		return nil, fmt.Errorf("error getting absolute path for labels yaml file %s: %w", labelsYamlFname, err)
+	}
+	labelFs := os.DirFS(filepath.Dir(path))
+	return getLabels(filepath.Base(path), labelFs.(fs.ReadFileFS))
+}
+
+func getLabels(fname string, labelsFs fs.ReadFileFS) ([]Label, error) {
+	// Read and parse the labels yaml file.
+	yamlBytes, err := labelsFs.ReadFile(fname)
+	if err != nil {
+		return nil, fmt.Errorf("error reading label yaml file %s", fname)
+	}
+	type yamlLabel struct {
+		Label
+		Rule string `yaml:"rule"`
+	}
+	yamlLabels := make(map[string]yamlLabel)
+	if err := yaml.Unmarshal(yamlBytes, &yamlLabels); err != nil {
+		return nil, fmt.Errorf("error unmarshalling labels yaml: %w", err)
+	}
+	labels := make([]Label, 0, len(yamlLabels))
+	// Read each label's classification rule.
+	var errs []error
+	for name, lbl := range yamlLabels {
+		// We assume that the rule will be defined with a relative path in the
+		// labels yaml file. If the rule is absolute, we create a new fs for
+		// the rule's directory. This implies that embedded labels will not work
+		// if the rule is absolute, but that should not be a problem because
+		// they should always be relative.
+		ruleFname := filepath.Join(filepath.Dir(fname), lbl.Rule)
+		ruleFs := labelsFs
+		if filepath.IsAbs(ruleFname) {
+			ruleFs = os.DirFS(filepath.Dir(ruleFname)).(fs.ReadFileFS)
+			ruleFname = filepath.Base(ruleFname)
+		}
+		rule, err := readLabelRule(ruleFname, ruleFs)
+		if err != nil {
+			errs = append(errs, fmt.Errorf("error reading classification rule for label %s: %w", name, err))
+			continue
+		}
+		lbl.Name = name
+		lbl.ClassificationRule = rule
+		labels = append(labels, lbl.Label)
+	}
+	if len(errs) > 0 {
+		return labels, InvalidLabelsError{errs}
+	}
+	return labels, nil
+}
+
+func readLabelRule(fname string, labelFs fs.ReadFileFS) (*ast.Module, error) {
+	b, err := labelFs.ReadFile(fname)
+	if err != nil {
+		return nil, fmt.Errorf("error reading rego file %s: %w", fname, err)
+	}
+	rule, err := parseRego(string(b))
+	if err != nil {
+		return nil, fmt.Errorf("error parsing classification rule for file %s: %w", fname, err)
+	}
+	return rule, nil
+}
+
+func parseRego(code string) (*ast.Module, error) {
+	log.Tracef("classifier module code: '%s'", code)
+	module, err := ast.ParseModule("classifier", code)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing rego code: %w", err)
+	}
+	return module, nil
+}