breakpilot-compliance/ai-compliance-sdk/internal/ucca/canonical_control_loader.go

package ucca

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

// CanonicalControl represents a single independently authored security control.
type CanonicalControl struct {
	ControlID            string             `json:"control_id"`
	Title                string             `json:"title"`
	Domain               string             `json:"domain"`
	Severity             string             `json:"severity"` // low, medium, high, critical
	RiskScore            float64            `json:"risk_score"`
	ImplementationEffort string             `json:"implementation_effort"` // s, m, l, xl
	Objective            string             `json:"objective"`
	Rationale            string             `json:"rationale"`
	Scope                CanonicalScope     `json:"scope"`
	Requirements         []string           `json:"requirements"`
	TestProcedure        []string           `json:"test_procedure"`
	Evidence             []CanonicalEvidence `json:"evidence"`
	OpenAnchors          []OpenAnchor       `json:"open_anchors"`
	Tags                 []string           `json:"tags"`
}

// CanonicalScope defines where a control applies.
type CanonicalScope struct {
	Platforms   []string `json:"platforms"`
	Components  []string `json:"components"`
	DataClasses []string `json:"data_classes"`
}

// CanonicalEvidence describes a required evidence item.
type CanonicalEvidence struct {
	Type        string `json:"type"`
	Description string `json:"description"`
}

// OpenAnchor links a control to an open-source framework reference.
type OpenAnchor struct {
	Framework string `json:"framework"`
	Ref       string `json:"ref"`
	URL       string `json:"url"`
}

// CanonicalDomain groups controls by security domain.
type CanonicalDomain struct {
	ID        string `json:"id"`
	Name      string `json:"name"`
	Objective string `json:"objective"`
}

// CanonicalFramework is the framework metadata.
type CanonicalFramework struct {
	ID          string `json:"id"`
	Name        string `json:"name"`
	Version     string `json:"version"`
	Description string `json:"description"`
}

// CanonicalControlLibrary is the top-level JSON structure.
type CanonicalControlLibrary struct {
	Version       string             `json:"version"`
	Schema        string             `json:"schema"`
	Generated     string             `json:"generated"`
	Framework     CanonicalFramework `json:"framework"`
	TotalControls int                `json:"total_controls"`
	Domains       []CanonicalDomain  `json:"domains"`
	Controls      []CanonicalControl `json:"controls"`
}

// CanonicalControlIndex provides fast lookup of canonical controls.
type CanonicalControlIndex struct {
	ByID          map[string]*CanonicalControl
	ByDomain      map[string][]*CanonicalControl
	BySeverity    map[string][]*CanonicalControl
	ByFramework   map[string][]*CanonicalControl // framework ref -> controls
	Domains       []CanonicalDomain
	Framework     CanonicalFramework
	AllControls   []*CanonicalControl
}

// LoadCanonicalControls loads the canonical control library from JSON.
func LoadCanonicalControls() (*CanonicalControlIndex, error) {
	data, err := readCanonicalControlsFile()
	if err != nil {
		return nil, err
	}

	var library CanonicalControlLibrary
	if err := json.Unmarshal(data, &library); err != nil {
		return nil, fmt.Errorf("failed to parse canonical controls: %w", err)
	}

	return buildCanonicalIndex(&library), nil
}

func readCanonicalControlsFile() ([]byte, error) {
	candidates := []string{
		"policies/canonical_controls_v1.json",
		"../policies/canonical_controls_v1.json",
		"../../policies/canonical_controls_v1.json",
	}

	_, filename, _, ok := runtime.Caller(0)
	if ok {
		srcDir := filepath.Dir(filename)
		candidates = append(candidates,
			filepath.Join(srcDir, "../../policies/canonical_controls_v1.json"),
		)
	}

	for _, p := range candidates {
		abs, err := filepath.Abs(p)
		if err != nil {
			continue
		}
		data, err := os.ReadFile(abs)
		if err == nil {
			return data, nil
		}
	}

	return nil, fmt.Errorf("canonical_controls_v1.json not found in any candidate path")
}

func buildCanonicalIndex(library *CanonicalControlLibrary) *CanonicalControlIndex {
	idx := &CanonicalControlIndex{
		ByID:        make(map[string]*CanonicalControl),
		ByDomain:    make(map[string][]*CanonicalControl),
		BySeverity:  make(map[string][]*CanonicalControl),
		ByFramework: make(map[string][]*CanonicalControl),
		Domains:     library.Domains,
		Framework:   library.Framework,
	}

	for i := range library.Controls {
		ctrl := &library.Controls[i]

		idx.ByID[ctrl.ControlID] = ctrl
		idx.ByDomain[ctrl.Domain] = append(idx.ByDomain[ctrl.Domain], ctrl)
		idx.BySeverity[ctrl.Severity] = append(idx.BySeverity[ctrl.Severity], ctrl)
		idx.AllControls = append(idx.AllControls, ctrl)

		for _, anchor := range ctrl.OpenAnchors {
			idx.ByFramework[anchor.Framework] = append(idx.ByFramework[anchor.Framework], ctrl)
		}
	}

	return idx
}

// GetControl returns a control by its ID (e.g. "AUTH-001").
func (idx *CanonicalControlIndex) GetControl(id string) (*CanonicalControl, bool) {
	ctrl, ok := idx.ByID[strings.ToUpper(id)]
	return ctrl, ok
}

// GetControlsByDomain returns all controls for a domain (e.g. "AUTH").
func (idx *CanonicalControlIndex) GetControlsByDomain(domain string) []*CanonicalControl {
	return idx.ByDomain[strings.ToUpper(domain)]
}

// GetControlsBySeverity returns all controls with a given severity.
func (idx *CanonicalControlIndex) GetControlsBySeverity(severity string) []*CanonicalControl {
	return idx.BySeverity[strings.ToLower(severity)]
}

// GetControlsByFramework returns all controls anchored to a framework (e.g. "OWASP ASVS").
func (idx *CanonicalControlIndex) GetControlsByFramework(framework string) []*CanonicalControl {
	return idx.ByFramework[framework]
}

// ValidateControlID checks if a control ID exists.
func (idx *CanonicalControlIndex) ValidateControlID(id string) bool {
	_, ok := idx.ByID[strings.ToUpper(id)]
	return ok
}