flex-auth/internal/policy/package.go

package policy

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"strings"

	"github.com/open-policy-agent/opa/ast"
	"github.com/open-policy-agent/opa/rego"
	"gopkg.in/yaml.v3"

	"github.com/netkingdom/flex-auth/pkg/api"
)

// Diagnostic is a validation message emitted while loading or evaluating a
// policy package.
type Diagnostic struct {
	Code     string         `json:"code"`
	Severity string         `json:"severity"`
	Message  string         `json:"message"`
	Fields   []string       `json:"fields,omitempty"`
	Metadata map[string]any `json:"metadata,omitempty"`
}

// CodeBlock is a fenced code block extracted from a policy package document.
type CodeBlock struct {
	Language  string   `json:"language"`
	Tags      []string `json:"tags,omitempty"`
	Info      string   `json:"info,omitempty"`
	Body      string   `json:"body"`
	StartLine int      `json:"start_line"`
}

// Package is a loaded Rego-in-Markdown policy package.
type Package struct {
	Source        string                    `json:"source,omitempty"`
	Metadata      api.PolicyPackageMetadata `json:"metadata"`
	Prose         string                    `json:"prose,omitempty"`
	RuleBlocks    []CodeBlock               `json:"rule_blocks,omitempty"`
	TestBlocks    []CodeBlock               `json:"test_blocks,omitempty"`
	FixtureBlocks []CodeBlock               `json:"fixture_blocks,omitempty"`
	RegoModule    string                    `json:"rego_module,omitempty"`
	TestModule    string                    `json:"test_module,omitempty"`
	TestPackage   string                    `json:"test_package,omitempty"`
	Fixtures      []api.PolicyFixture       `json:"fixtures,omitempty"`
	Validation    ValidationResult          `json:"validation,omitempty"`
	Valid         bool                      `json:"valid"`
}

// ValidationResult captures OPA and CARING validation outcomes for a package.
type ValidationResult struct {
	Valid          bool                           `json:"valid"`
	Diagnostics    []Diagnostic                   `json:"diagnostics,omitempty"`
	CaringFindings []api.CaringConformanceFinding `json:"caring_findings,omitempty"`
	Tests          []TestResult                   `json:"tests,omitempty"`
	Fixtures       []FixtureResult                `json:"fixtures,omitempty"`
}

// TestResult captures one OPA test rule result.
type TestResult struct {
	Name   string `json:"name"`
	Passed bool   `json:"passed"`
	Error  string `json:"error,omitempty"`
}

// FixtureResult captures one embedded or referenced fixture evaluation result.
type FixtureResult struct {
	ID       string                  `json:"id,omitempty"`
	Passed   bool                    `json:"passed"`
	Expected api.DecisionExpectation `json:"expected"`
	Actual   api.DecisionExpectation `json:"actual"`
	Error    string                  `json:"error,omitempty"`
}

// LoadFile loads a policy package Markdown document from disk.
func LoadFile(path string) (*Package, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read policy package: %w", err)
	}

	pkg, err := Load(data, path)
	if err != nil {
		return nil, err
	}
	if err := pkg.loadExternalFixtures(filepath.Dir(path)); err != nil {
		return nil, err
	}
	return pkg, nil
}

// Load parses a policy package Markdown document.
func Load(data []byte, source string) (*Package, error) {
	frontmatter, body, err := splitFrontmatter(string(data))
	if err != nil {
		return nil, err
	}

	var metadata api.PolicyPackageMetadata
	if err := yaml.Unmarshal([]byte(frontmatter), &metadata); err != nil {
		return nil, fmt.Errorf("unmarshal policy package frontmatter: %w", err)
	}

	extracted, err := extractMarkdown(body)
	if err != nil {
		return nil, err
	}

	pkg := &Package{
		Source:        source,
		Metadata:      metadata,
		Prose:         strings.TrimSpace(extracted.prose),
		RuleBlocks:    extracted.ruleBlocks,
		TestBlocks:    extracted.testBlocks,
		FixtureBlocks: extracted.fixtureBlocks,
	}

	pkg.RegoModule, _ = normalizeRuleModule(metadata.Package, pkg.RuleBlocks)
	pkg.TestModule, pkg.TestPackage, _ = normalizeTestModule(metadata.Package, pkg.TestBlocks)

	for _, block := range pkg.FixtureBlocks {
		fixtures, err := parseFixtureYAML(block.Body)
		if err != nil {
			return nil, fmt.Errorf("parse fixture block at line %d: %w", block.StartLine, err)
		}
		pkg.Fixtures = append(pkg.Fixtures, fixtures...)
	}

	return pkg, nil
}

// LoadAndValidateFile loads a policy package and immediately validates it.
func LoadAndValidateFile(ctx context.Context, path string) (*Package, error) {
	pkg, err := LoadFile(path)
	if err != nil {
		return nil, err
	}
	pkg.Validate(ctx)
	return pkg, nil
}

// Validate runs metadata, CARING, OPA parse/test, and fixture validation.
func (p *Package) Validate(ctx context.Context) ValidationResult {
	result := ValidationResult{}

	result.Diagnostics = append(result.Diagnostics, p.metadataDiagnostics()...)
	result.CaringFindings = append(result.CaringFindings, p.caringFindings()...)
	if len(p.Fixtures) == 0 {
		result.Diagnostics = append(result.Diagnostics, Diagnostic{
			Code:     "POLICY-FIXTURE-MISSING",
			Severity: "error",
			Message:  "policy package must include at least one embedded or referenced fixture",
		})
	}

	regoModule, regoDiagnostics := normalizeRuleModule(p.Metadata.Package, p.RuleBlocks)
	result.Diagnostics = append(result.Diagnostics, regoDiagnostics...)
	p.RegoModule = regoModule

	testModule, testPackage, testDiagnostics := normalizeTestModule(p.Metadata.Package, p.TestBlocks)
	result.Diagnostics = append(result.Diagnostics, testDiagnostics...)
	p.TestModule = testModule
	p.TestPackage = testPackage

	parseOK := true
	if p.RegoModule != "" {
		if _, err := ast.ParseModule(moduleFilename(p.Source, "rego"), p.RegoModule); err != nil {
			parseOK = false
			result.Diagnostics = append(result.Diagnostics, Diagnostic{
				Code:     "OPA-PARSE",
				Severity: "error",
				Message:  err.Error(),
			})
		}
	}
	if p.TestModule != "" {
		if _, err := ast.ParseModule(moduleFilename(p.Source, "test.rego"), p.TestModule); err != nil {
			parseOK = false
			result.Diagnostics = append(result.Diagnostics, Diagnostic{
				Code:     "OPA-TEST-PARSE",
				Severity: "error",
				Message:  err.Error(),
			})
		}
	}

	if parseOK && p.RegoModule != "" {
		result.Tests = p.runTests(ctx)
		result.Fixtures = p.runFixtures(ctx)
	}

	result.Valid = validationPassed(result, p.Metadata.Caring.Enforce)
	p.Validation = result
	p.Valid = result.Valid
	return result
}

func (p *Package) loadExternalFixtures(baseDir string) error {
	for _, fixturePath := range p.Metadata.Fixtures {
		path := fixturePath
		if !filepath.IsAbs(path) {
			path = filepath.Join(baseDir, fixturePath)
		}

		data, err := os.ReadFile(path)
		if err != nil {
			return fmt.Errorf("read policy fixture %q: %w", fixturePath, err)
		}
		fixtures, err := parseFixtureYAML(string(data))
		if err != nil {
			return fmt.Errorf("parse policy fixture %q: %w", fixturePath, err)
		}
		p.Fixtures = append(p.Fixtures, fixtures...)
	}
	return nil
}

func (p *Package) metadataDiagnostics() []Diagnostic {
	var diagnostics []Diagnostic
	if p.Metadata.ID == "" {
		diagnostics = append(diagnostics, requiredDiagnostic("POLICY-METADATA-ID", "id", "policy package id is required"))
	}
	if p.Metadata.Version == "" {
		diagnostics = append(diagnostics, requiredDiagnostic("POLICY-METADATA-VERSION", "version", "policy package version is required"))
	}
	if p.Metadata.Package == "" {
		diagnostics = append(diagnostics, requiredDiagnostic("POLICY-METADATA-PACKAGE", "package", "OPA package path is required"))
	}
	return diagnostics
}

func (p *Package) caringFindings() []api.CaringConformanceFinding {
	var findings []api.CaringConformanceFinding
	caring := p.Metadata.Caring
	if caring.Profile == "" {
		findings = append(findings, caringFinding("CARING-POLICY-PROFILE", "error", "policy package must declare a CARING profile", "caring.profile"))
	} else if caring.Profile != api.CaringProfileCaring040RC2 {
		findings = append(findings, caringFinding("CARING-POLICY-PROFILE", "error", fmt.Sprintf("unsupported CARING profile %q", caring.Profile), "caring.profile"))
	}

	addMissing := func(empty bool, field, label string) {
		if empty {
			findings = append(findings, caringFinding("CARING-POLICY-MISSING-DIMENSION", "warning", "policy package should declare governed "+label, field))
		}
	}
	addMissing(len(caring.CanonicalRoles) == 0, "caring.canonical_roles", "canonical roles")
	addMissing(len(caring.OrganizationRelations) == 0, "caring.organization_relations", "organization relations")
	addMissing(len(caring.Scopes) == 0, "caring.scopes", "scopes")
	addMissing(len(caring.Planes) == 0, "caring.planes", "planes")
	addMissing(len(caring.Capabilities) == 0, "caring.capabilities", "capabilities")
	addMissing(len(caring.ExposureModes) == 0, "caring.exposure_modes", "exposure modes")
	addMissing(len(caring.Conditions) == 0, "caring.conditions", "conditions")
	addMissing(len(caring.Restrictions) == 0, "caring.restrictions", "restrictions")
	return findings
}

func requiredDiagnostic(code, field, message string) Diagnostic {
	return Diagnostic{
		Code:     code,
		Severity: "error",
		Message:  message,
		Fields:   []string{field},
	}
}

func caringFinding(code, severity, message, field string) api.CaringConformanceFinding {
	return api.CaringConformanceFinding{
		Code:     code,
		Severity: severity,
		Message:  message,
		Fields:   []string{field},
	}
}

func (p *Package) runTests(ctx context.Context) []TestResult {
	if p.TestModule == "" {
		return nil
	}

	module, err := ast.ParseModule(moduleFilename(p.Source, "test.rego"), p.TestModule)
	if err != nil {
		return []TestResult{{Name: "parse", Passed: false, Error: err.Error()}}
	}

	var names []string
	for _, rule := range module.Rules {
		name := string(rule.Head.Name)
		if strings.HasPrefix(name, "test_") {
			names = append(names, name)
		}
	}
	sort.Strings(names)
	if len(names) == 0 {
		return []TestResult{{
			Name:   "test_*",
			Passed: false,
			Error:  "test module does not contain any test_* rules",
		}}
	}

	results := make([]TestResult, 0, len(names))
	for _, name := range names {
		query := "data." + p.TestPackage + "." + name
		passed, err := p.evalBool(ctx, query, nil, true)
		result := TestResult{Name: name, Passed: passed}
		if err != nil {
			result.Error = err.Error()
		} else if !passed {
			result.Error = "test rule evaluated to false or undefined"
		}
		results = append(results, result)
	}
	return results
}

func (p *Package) runFixtures(ctx context.Context) []FixtureResult {
	results := make([]FixtureResult, 0, len(p.Fixtures))
	for _, fixture := range p.Fixtures {
		actual, err := p.evaluateDecision(ctx, fixture.Request)
		result := FixtureResult{
			ID:       fixture.ID,
			Expected: fixture.Expect,
			Actual:   actual,
		}
		if err != nil {
			result.Error = err.Error()
		} else if expectationMatches(fixture.Expect, actual) {
			result.Passed = true
		} else {
			result.Error = "fixture expectation did not match actual decision"
		}
		results = append(results, result)
	}
	return results
}

func (p *Package) evaluateDecision(ctx context.Context, request api.CheckRequest) (api.DecisionExpectation, error) {
	input, err := toRegoInput(request)
	if err != nil {
		return api.DecisionExpectation{}, err
	}

	query := "data." + p.Metadata.Package + ".decision"
	results, err := p.eval(ctx, query, input, false)
	if err != nil {
		return api.DecisionExpectation{}, err
	}
	if len(results) == 0 || len(results[0].Expressions) == 0 {
		return api.DecisionExpectation{}, fmt.Errorf("decision query %q was undefined", query)
	}

	data, err := json.Marshal(results[0].Expressions[0].Value)
	if err != nil {
		return api.DecisionExpectation{}, fmt.Errorf("marshal decision result: %w", err)
	}

	var decision api.DecisionExpectation
	if err := json.Unmarshal(data, &decision); err != nil {
		return api.DecisionExpectation{}, fmt.Errorf("unmarshal decision result: %w", err)
	}
	return decision, nil
}

func (p *Package) evalBool(ctx context.Context, query string, input map[string]any, includeTests bool) (bool, error) {
	results, err := p.eval(ctx, query, input, includeTests)
	if err != nil {
		return false, err
	}
	for _, result := range results {
		if len(result.Expressions) == 0 {
			return true, nil
		}
		if passed, ok := result.Expressions[0].Value.(bool); ok && passed {
			return true, nil
		}
	}
	return false, nil
}

func (p *Package) eval(ctx context.Context, query string, input map[string]any, includeTests bool) (rego.ResultSet, error) {
	options := []func(*rego.Rego){
		rego.Query(query),
		rego.Module(moduleFilename(p.Source, "rego"), p.RegoModule),
	}
	if includeTests && p.TestModule != "" {
		options = append(options, rego.Module(moduleFilename(p.Source, "test.rego"), p.TestModule))
	}
	if input != nil {
		options = append(options, rego.Input(input))
	}

	prepared, err := rego.New(options...).PrepareForEval(ctx)
	if err != nil {
		return nil, err
	}
	return prepared.Eval(ctx)
}

func expectationMatches(expected, actual api.DecisionExpectation) bool {
	if expected.Effect != "" && expected.Effect != actual.Effect {
		return false
	}
	if expected.Reason != "" && expected.Reason != actual.Reason {
		return false
	}
	if len(expected.Obligations) > 0 && !reflect.DeepEqual(expected.Obligations, actual.Obligations) {
		return false
	}
	if len(expected.ConformanceFindings) > 0 && !reflect.DeepEqual(expected.ConformanceFindings, actual.ConformanceFindings) {
		return false
	}
	return true
}

func toRegoInput(value any) (map[string]any, error) {
	data, err := json.Marshal(value)
	if err != nil {
		return nil, fmt.Errorf("marshal input: %w", err)
	}
	var out map[string]any
	if err := json.Unmarshal(data, &out); err != nil {
		return nil, fmt.Errorf("unmarshal input: %w", err)
	}
	return out, nil
}

func validationPassed(result ValidationResult, enforceCaring bool) bool {
	for _, diagnostic := range result.Diagnostics {
		if diagnostic.Severity == "error" {
			return false
		}
	}
	for _, finding := range result.CaringFindings {
		if finding.Severity == "error" || (enforceCaring && finding.Severity != "info") {
			return false
		}
	}
	for _, test := range result.Tests {
		if !test.Passed {
			return false
		}
	}
	for _, fixture := range result.Fixtures {
		if !fixture.Passed {
			return false
		}
	}
	return true
}

func normalizeRuleModule(packageName string, blocks []CodeBlock) (string, []Diagnostic) {
	module, _, diagnostics := normalizeModule(packageName, "", blocks, "POLICY-REGO")
	return module, diagnostics
}

func normalizeTestModule(packageName string, blocks []CodeBlock) (string, string, []Diagnostic) {
	if len(blocks) == 0 {
		return "", "", []Diagnostic{{
			Code:     "POLICY-REGO-TEST-MISSING",
			Severity: "error",
			Message:  "policy package must include at least one rego test block",
		}}
	}

	defaultPackage := ""
	if packageName != "" {
		defaultPackage = packageName + "_test"
	}
	module, testPackage, diagnostics := normalizeModule(defaultPackage, "test", blocks, "POLICY-REGO-TEST")
	return module, testPackage, diagnostics
}

func normalizeModule(defaultPackage, label string, blocks []CodeBlock, codePrefix string) (string, string, []Diagnostic) {
	if len(blocks) == 0 {
		return "", "", []Diagnostic{{
			Code:     codePrefix + "-MISSING",
			Severity: "error",
			Message:  "policy package must include at least one " + labelOrDefault(label, "rule") + " block",
		}}
	}

	var packageName string
	var parts []string
	var diagnostics []Diagnostic
	for _, block := range blocks {
		body, declared := stripPackageDeclaration(block.Body)
		if declared != "" {
			if packageName == "" {
				packageName = declared
			} else if declared != packageName {
				diagnostics = append(diagnostics, Diagnostic{
					Code:     codePrefix + "-PACKAGE-MISMATCH",
					Severity: "error",
					Message:  fmt.Sprintf("block at line %d declares package %q, expected %q", block.StartLine, declared, packageName),
					Fields:   []string{"package"},
				})
			}
		}
		if strings.TrimSpace(body) != "" {
			parts = append(parts, strings.TrimSpace(body))
		}
	}

	if defaultPackage != "" {
		if packageName != "" && packageName != defaultPackage {
			diagnostics = append(diagnostics, Diagnostic{
				Code:     codePrefix + "-PACKAGE-MISMATCH",
				Severity: "error",
				Message:  fmt.Sprintf("declared package %q does not match frontmatter package %q", packageName, defaultPackage),
				Fields:   []string{"package"},
			})
		}
		packageName = defaultPackage
	}
	if packageName == "" {
		diagnostics = append(diagnostics, Diagnostic{
			Code:     codePrefix + "-PACKAGE-MISSING",
			Severity: "error",
			Message:  "policy package cannot build a Rego module without a package name",
			Fields:   []string{"package"},
		})
		return "", "", diagnostics
	}

	module := "package " + packageName + "\n\n" + strings.Join(parts, "\n\n")
	return strings.TrimSpace(module) + "\n", packageName, diagnostics
}

func labelOrDefault(label, fallback string) string {
	if label == "" {
		return fallback
	}
	return label
}

func stripPackageDeclaration(body string) (string, string) {
	lines := strings.Split(body, "\n")
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || strings.HasPrefix(trimmed, "#") {
			continue
		}
		if strings.HasPrefix(trimmed, "package ") {
			declared := strings.TrimSpace(strings.TrimPrefix(trimmed, "package "))
			lines = append(lines[:i], lines[i+1:]...)
			return strings.Join(lines, "\n"), declared
		}
		return body, ""
	}
	return body, ""
}

func parseFixtureYAML(body string) ([]api.PolicyFixture, error) {
	var node yaml.Node
	if err := yaml.Unmarshal([]byte(body), &node); err != nil {
		return nil, err
	}
	if len(node.Content) == 0 {
		return nil, nil
	}

	root := node.Content[0]
	if root.Kind == yaml.SequenceNode {
		var fixtures []api.PolicyFixture
		if err := root.Decode(&fixtures); err != nil {
			return nil, err
		}
		return fixtures, nil
	}

	var fixture api.PolicyFixture
	if err := root.Decode(&fixture); err != nil {
		return nil, err
	}
	return []api.PolicyFixture{fixture}, nil
}

type markdownExtraction struct {
	prose         string
	ruleBlocks    []CodeBlock
	testBlocks    []CodeBlock
	fixtureBlocks []CodeBlock
}

func extractMarkdown(body string) (markdownExtraction, error) {
	var out markdownExtraction
	var prose strings.Builder
	var block strings.Builder
	var current *CodeBlock

	lines := strings.Split(body, "\n")
	for i, line := range lines {
		lineNo := i + 1
		trimmed := strings.TrimSpace(line)
		if current == nil {
			if strings.HasPrefix(trimmed, "```") {
				info := strings.TrimSpace(strings.TrimPrefix(trimmed, "```"))
				current = &CodeBlock{
					Info:      info,
					StartLine: lineNo,
				}
				current.Language, current.Tags = parseFenceInfo(info)
				block.Reset()
				continue
			}
			prose.WriteString(line)
			if i < len(lines)-1 {
				prose.WriteByte('\n')
			}
			continue
		}

		if strings.HasPrefix(trimmed, "```") {
			current.Body = strings.TrimRight(block.String(), "\n")
			switch {
			case current.Language == "rego" && hasTag(current.Tags, "test"):
				out.testBlocks = append(out.testBlocks, *current)
			case current.Language == "rego":
				out.ruleBlocks = append(out.ruleBlocks, *current)
			case (current.Language == "yaml" || current.Language == "yml") && hasTag(current.Tags, "fixture"):
				out.fixtureBlocks = append(out.fixtureBlocks, *current)
			}
			current = nil
			continue
		}

		block.WriteString(line)
		block.WriteByte('\n')
	}
	if current != nil {
		return out, fmt.Errorf("unterminated fenced code block starting at line %d", current.StartLine)
	}
	out.prose = prose.String()
	return out, nil
}

func parseFenceInfo(info string) (string, []string) {
	fields := strings.Fields(info)
	if len(fields) == 0 {
		return "", nil
	}
	language := strings.ToLower(fields[0])
	tags := make([]string, 0, len(fields)-1)
	for _, field := range fields[1:] {
		tags = append(tags, strings.ToLower(field))
	}
	return language, tags
}

func hasTag(tags []string, want string) bool {
	for _, tag := range tags {
		if tag == want {
			return true
		}
	}
	return false
}

func splitFrontmatter(document string) (string, string, error) {
	document = strings.TrimPrefix(document, "\ufeff")
	lines := strings.SplitAfter(document, "\n")
	if len(lines) == 0 || strings.TrimSpace(lines[0]) != "---" {
		return "", "", fmt.Errorf("policy package must start with YAML frontmatter")
	}

	for i := 1; i < len(lines); i++ {
		if strings.TrimSpace(lines[i]) == "---" {
			return strings.Join(lines[1:i], ""), strings.Join(lines[i+1:], ""), nil
		}
	}
	return "", "", fmt.Errorf("policy package frontmatter is not closed")
}

func moduleFilename(source, suffix string) string {
	if source == "" {
		return "policy." + suffix
	}
	return source + "." + suffix
}