github.com/datreeio/datree@v1.9.22-rc/pkg/jsonSchemaValidator/extensions/customKeyRegoDefinition.go (about)

     1  // This file defines a custom key to implement the logic for rego rule:
     2  
     3  package jsonSchemaValidator
     4  
     5  import (
     6  	"context"
     7  	"encoding/json"
     8  	"fmt"
     9  	"strings"
    10  
    11  	"github.com/open-policy-agent/opa/rego"
    12  	"github.com/santhosh-tekuri/jsonschema/v5"
    13  )
    14  
    15  const RegoDefinitionCustomKey = "regoDefinition"
    16  
    17  type CustomKeyRegoDefinitionCompiler struct{}
    18  
    19  type CustomKeyRegoDefinitionSchema map[string]interface{}
    20  
    21  var CustomKeyRegoRule = jsonschema.MustCompileString("customKeyRegoDefinition.json", `{
    22  	"properties" : {
    23  		"regoDefinition": {
    24  			"type": "object"
    25  		}
    26  	}
    27  }`)
    28  
    29  func (CustomKeyRegoDefinitionCompiler) Compile(ctx jsonschema.CompilerContext, m map[string]interface{}) (jsonschema.ExtSchema, error) {
    30  	if customKeyRegoRule, ok := m[RegoDefinitionCustomKey]; ok {
    31  		customKeyRegoRuleObj, validObject := customKeyRegoRule.(map[string]interface{})
    32  		if !validObject {
    33  			return nil, fmt.Errorf("regoDefinition must be an object")
    34  		}
    35  
    36  		regoDefinitionSchema, err := convertCustomKeyRegoDefinitionSchemaToRegoDefinitionSchema(customKeyRegoRuleObj)
    37  		if err != nil {
    38  			return nil, err
    39  		}
    40  
    41  		if regoDefinitionSchema.Code == "" {
    42  			return nil, fmt.Errorf("regoDefinition.code can't be empty")
    43  		}
    44  
    45  		return CustomKeyRegoDefinitionSchema(customKeyRegoRuleObj), nil
    46  	}
    47  	return nil, nil
    48  }
    49  
    50  type RegoDefinition struct {
    51  	Libs []string `json:"libs"`
    52  	Code string   `json:"code"`
    53  }
    54  
    55  func (customKeyRegoDefinitionSchema CustomKeyRegoDefinitionSchema) Validate(ctx jsonschema.ValidationContext, dataValue interface{}) error {
    56  	regoDefinitionSchema, err := convertCustomKeyRegoDefinitionSchemaToRegoDefinitionSchema(customKeyRegoDefinitionSchema)
    57  	if err != nil {
    58  		return ctx.Error(CustomKeyValidationErrorKeyPath, err.Error())
    59  	}
    60  
    61  	regoCtx := context.Background()
    62  
    63  	regoObject, err := retrieveRegoFromSchema(regoDefinitionSchema)
    64  	if err != nil {
    65  		return ctx.Error(CustomKeyValidationErrorKeyPath, "can't compile rego code, %s", err.Error())
    66  	}
    67  
    68  	// Create a prepared query that can be evaluated.
    69  	query, err := regoObject.PrepareForEval(regoCtx)
    70  	if err != nil {
    71  		return ctx.Error(CustomKeyValidationErrorKeyPath, "can't compile rego code, %s", err.Error())
    72  	}
    73  
    74  	// Execute the prepared query.
    75  	rs, err := query.Eval(regoCtx, rego.EvalInput(dataValue))
    76  
    77  	if err != nil {
    78  		return ctx.Error(CustomKeyValidationErrorKeyPath, "failed to evaluate rego due to %s", err.Error())
    79  	}
    80  
    81  	if len(rs) != 1 || len(rs[0].Expressions) != 1 {
    82  		return ctx.Error(CustomKeyValidationErrorKeyPath, "failed to evaluate rego, unexpected results")
    83  	}
    84  
    85  	resultValues := (rs[0].Expressions[0].Value).([]interface{})
    86  	for _, resultValue := range resultValues {
    87  		violationReturnValue, ok := resultValue.(bool)
    88  		if !ok {
    89  			return ctx.Error(CustomKeyValidationErrorKeyPath, "violation needs to return a boolean")
    90  		}
    91  		if violationReturnValue {
    92  			return ctx.Error(RegoDefinitionCustomKey, "values in data value %v do not match", dataValue)
    93  		}
    94  	}
    95  
    96  	return nil
    97  }
    98  
    99  func getPackageFromRegoCode(regoCode string) (string, error) {
   100  	const PACKAGE = "package"
   101  	// find the index of string "package"
   102  	index := strings.Index(regoCode, PACKAGE)
   103  	if index == -1 {
   104  		return "", fmt.Errorf("rego code must have a package")
   105  	}
   106  	// get next single word after "package"
   107  	packageStr := strings.Fields(regoCode[index:])
   108  	return packageStr[1], nil
   109  }
   110  
   111  func retrieveRegoFromSchema(regoDefinitionSchema *RegoDefinition) (*rego.Rego, error) {
   112  	const mainModuleFileName = "main.rego"
   113  	const regoFunctionEntryPoint = "violation"
   114  
   115  	mainRegoPackage, err := getPackageFromRegoCode(regoDefinitionSchema.Code)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  
   120  	var regoObjectParts []func(r *rego.Rego)
   121  	regoObjectParts = append(regoObjectParts, rego.Query("data."+mainRegoPackage+"."+regoFunctionEntryPoint))
   122  
   123  	regoObjectParts = append(regoObjectParts, rego.Module(mainModuleFileName, regoDefinitionSchema.Code))
   124  
   125  	for _, lib := range regoDefinitionSchema.Libs {
   126  		libPackageName, err := getPackageFromRegoCode(lib)
   127  		if err != nil {
   128  			return nil, err
   129  		}
   130  		regoObjectParts = append(regoObjectParts, rego.Module(libPackageName, lib))
   131  	}
   132  	regoObject := rego.New(regoObjectParts...)
   133  	return regoObject, nil
   134  }
   135  
   136  func convertCustomKeyRegoDefinitionSchemaToRegoDefinitionSchema(regoDefinitionSchema CustomKeyRegoDefinitionSchema) (*RegoDefinition, error) {
   137  	b, err := json.Marshal(regoDefinitionSchema)
   138  	if err != nil {
   139  		return nil, fmt.Errorf("regoDefinition failed to marshal to json, %s", err.Error())
   140  	}
   141  
   142  	var regoDefinition RegoDefinition
   143  	err = json.Unmarshal(b, &regoDefinition)
   144  	if err != nil {
   145  		return nil, fmt.Errorf("regoDefinition must be an object of type RegoDefinition %s", err.Error())
   146  	}
   147  	return &regoDefinition, nil
   148  }