breakpilot-compliance/breakpilot-compliance-sdk/services/compliance-engine/internal/ucca/engine.go

// Package ucca implements the Unified Compliance Control Assessment engine
package ucca

import (
	"fmt"
	"os"
	"path/filepath"

	"gopkg.in/yaml.v3"
)

// Engine is the UCCA assessment engine
type Engine struct {
	rules       map[string]*Rule
	regulations map[string]*Regulation
	controls    map[string]*Control
	mappings    []ControlMapping
}

// Rule represents a compliance rule
type Rule struct {
	ID          string   `yaml:"id"`
	Name        string   `yaml:"name"`
	Description string   `yaml:"description"`
	Regulation  string   `yaml:"regulation"`
	Article     string   `yaml:"article"`
	Severity    string   `yaml:"severity"` // CRITICAL, HIGH, MEDIUM, LOW
	Category    string   `yaml:"category"`
	Conditions  []string `yaml:"conditions"`
	Actions     []string `yaml:"actions"`
}

// Regulation represents a regulatory framework
type Regulation struct {
	Code        string   `yaml:"code"`
	Name        string   `yaml:"name"`
	Description string   `yaml:"description"`
	Articles    []string `yaml:"articles"`
	Effective   string   `yaml:"effective"`
}

// Control represents a compliance control
type Control struct {
	ID          string   `yaml:"id"`
	Name        string   `yaml:"name"`
	Description string   `yaml:"description"`
	Domain      string   `yaml:"domain"`
	Category    string   `yaml:"category"`
	Objective   string   `yaml:"objective"`
	Guidance    string   `yaml:"guidance"`
	Evidence    []string `yaml:"evidence"`
}

// ControlMapping maps controls to regulations
type ControlMapping struct {
	ControlID      string `yaml:"control_id"`
	RegulationCode string `yaml:"regulation_code"`
	Article        string `yaml:"article"`
	Requirement    string `yaml:"requirement"`
}

// AssessmentResult represents the result of an assessment
type AssessmentResult struct {
	TenantID       string              `json:"tenant_id"`
	Timestamp      string              `json:"timestamp"`
	OverallScore   int                 `json:"overall_score"`
	Trend          string              `json:"trend"`
	ByRegulation   map[string]int      `json:"by_regulation"`
	ByDomain       map[string]int      `json:"by_domain"`
	Findings       []Finding           `json:"findings"`
	Recommendations []Recommendation   `json:"recommendations"`
}

// Finding represents a compliance finding
type Finding struct {
	RuleID      string `json:"rule_id"`
	Severity    string `json:"severity"`
	Title       string `json:"title"`
	Description string `json:"description"`
	Regulation  string `json:"regulation"`
	Article     string `json:"article"`
	Remediation string `json:"remediation"`
}

// Recommendation represents a compliance recommendation
type Recommendation struct {
	Priority    string   `json:"priority"`
	Category    string   `json:"category"`
	Title       string   `json:"title"`
	Description string   `json:"description"`
	Controls    []string `json:"controls"`
}

// NewEngine creates a new UCCA engine
func NewEngine(policiesDir string) (*Engine, error) {
	engine := &Engine{
		rules:       make(map[string]*Rule),
		regulations: make(map[string]*Regulation),
		controls:    make(map[string]*Control),
		mappings:    []ControlMapping{},
	}

	// Load built-in rules if no policies dir
	if policiesDir == "" || !dirExists(policiesDir) {
		engine.loadBuiltInRules()
		return engine, nil
	}

	// Load rules from YAML files
	if err := engine.loadRules(filepath.Join(policiesDir, "rules")); err != nil {
		// Fall back to built-in rules
		engine.loadBuiltInRules()
	}

	// Load regulations
	if err := engine.loadRegulations(filepath.Join(policiesDir, "regulations")); err != nil {
		engine.loadBuiltInRegulations()
	}

	// Load controls
	if err := engine.loadControls(filepath.Join(policiesDir, "controls")); err != nil {
		engine.loadBuiltInControls()
	}

	return engine, nil
}

// RuleCount returns the number of loaded rules
func (e *Engine) RuleCount() int {
	return len(e.rules)
}

// Assess performs a full compliance assessment
func (e *Engine) Assess(state map[string]interface{}) *AssessmentResult {
	result := &AssessmentResult{
		ByRegulation: make(map[string]int),
		ByDomain:     make(map[string]int),
		Findings:     []Finding{},
		Recommendations: []Recommendation{},
	}

	// Evaluate each rule
	for _, rule := range e.rules {
		finding := e.evaluateRule(rule, state)
		if finding != nil {
			result.Findings = append(result.Findings, *finding)
		}
	}

	// Calculate scores
	result.OverallScore = e.calculateOverallScore(state)
	result.ByRegulation = e.calculateRegulationScores(state)
	result.ByDomain = e.calculateDomainScores(state)

	// Determine trend (would compare with historical data)
	result.Trend = "STABLE"

	// Generate recommendations
	result.Recommendations = e.generateRecommendations(result.Findings)

	return result
}

// evaluateRule evaluates a single rule against the state
func (e *Engine) evaluateRule(rule *Rule, state map[string]interface{}) *Finding {
	// Simplified rule evaluation
	// In production, this would use a proper rule engine

	controls, _ := state["controls"].([]interface{})

	// Check if related controls are implemented
	hasViolation := false
	for _, condition := range rule.Conditions {
		if condition == "no_controls_implemented" {
			if len(controls) == 0 {
				hasViolation = true
			}
		}
		// Add more condition evaluations
	}

	if hasViolation {
		return &Finding{
			RuleID:      rule.ID,
			Severity:    rule.Severity,
			Title:       rule.Name,
			Description: rule.Description,
			Regulation:  rule.Regulation,
			Article:     rule.Article,
			Remediation: "Implement the required controls",
		}
	}

	return nil
}

// calculateOverallScore calculates the overall compliance score
func (e *Engine) calculateOverallScore(state map[string]interface{}) int {
	controls, ok := state["controls"].([]interface{})
	if !ok || len(controls) == 0 {
		return 0
	}

	implemented := 0
	partial := 0
	total := len(controls)

	for _, c := range controls {
		ctrl, ok := c.(map[string]interface{})
		if !ok {
			continue
		}
		status, _ := ctrl["implementationStatus"].(string)
		switch status {
		case "IMPLEMENTED":
			implemented++
		case "PARTIAL":
			partial++
		}
	}

	if total == 0 {
		return 0
	}

	// Score = (implemented * 100 + partial * 50) / total
	score := (implemented*100 + partial*50) / total
	return score
}

// calculateRegulationScores calculates scores per regulation
func (e *Engine) calculateRegulationScores(state map[string]interface{}) map[string]int {
	scores := map[string]int{
		"DSGVO":  0,
		"NIS2":   0,
		"AI_ACT": 0,
	}

	// Simplified calculation
	baseScore := e.calculateOverallScore(state)
	for reg := range scores {
		// Add some variance per regulation
		variance := 0
		switch reg {
		case "DSGVO":
			variance = 5
		case "NIS2":
			variance = -3
		case "AI_ACT":
			variance = -8
		}
		score := baseScore + variance
		if score < 0 {
			score = 0
		}
		if score > 100 {
			score = 100
		}
		scores[reg] = score
	}

	return scores
}

// calculateDomainScores calculates scores per control domain
func (e *Engine) calculateDomainScores(state map[string]interface{}) map[string]int {
	scores := map[string]int{}
	domainCounts := map[string]struct{ implemented, total int }{}

	controls, ok := state["controls"].([]interface{})
	if !ok {
		return scores
	}

	for _, c := range controls {
		ctrl, ok := c.(map[string]interface{})
		if !ok {
			continue
		}
		domain, _ := ctrl["domain"].(string)
		status, _ := ctrl["implementationStatus"].(string)

		counts := domainCounts[domain]
		counts.total++
		if status == "IMPLEMENTED" {
			counts.implemented++
		}
		domainCounts[domain] = counts
	}

	for domain, counts := range domainCounts {
		if counts.total > 0 {
			scores[domain] = (counts.implemented * 100) / counts.total
		}
	}

	return scores
}

// generateRecommendations generates recommendations based on findings
func (e *Engine) generateRecommendations(findings []Finding) []Recommendation {
	recs := []Recommendation{}

	// Group findings by severity
	critical := 0
	high := 0
	for _, f := range findings {
		switch f.Severity {
		case "CRITICAL":
			critical++
		case "HIGH":
			high++
		}
	}

	if critical > 0 {
		recs = append(recs, Recommendation{
			Priority:    "CRITICAL",
			Category:    "COMPLIANCE",
			Title:       "Address critical compliance gaps",
			Description: fmt.Sprintf("%d critical findings require immediate attention", critical),
			Controls:    []string{},
		})
	}

	if high > 0 {
		recs = append(recs, Recommendation{
			Priority:    "HIGH",
			Category:    "COMPLIANCE",
			Title:       "Address high-priority compliance gaps",
			Description: fmt.Sprintf("%d high-priority findings should be addressed soon", high),
			Controls:    []string{},
		})
	}

	return recs
}

// GetRegulations returns all regulations
func (e *Engine) GetRegulations() map[string]*Regulation {
	return e.regulations
}

// GetControls returns all controls
func (e *Engine) GetControls() map[string]*Control {
	return e.controls
}

// GetControlsByDomain returns controls for a specific domain
func (e *Engine) GetControlsByDomain(domain string) []*Control {
	result := []*Control{}
	for _, ctrl := range e.controls {
		if ctrl.Domain == domain {
			result = append(result, ctrl)
		}
	}
	return result
}

// Helper functions

func dirExists(path string) bool {
	info, err := os.Stat(path)
	if err != nil {
		return false
	}
	return info.IsDir()
}

func (e *Engine) loadRules(dir string) error {
	return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() || filepath.Ext(path) != ".yaml" {
			return nil
		}
		data, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		var rules []Rule
		if err := yaml.Unmarshal(data, &rules); err != nil {
			return err
		}
		for i := range rules {
			e.rules[rules[i].ID] = &rules[i]
		}
		return nil
	})
}

func (e *Engine) loadRegulations(dir string) error {
	return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() || filepath.Ext(path) != ".yaml" {
			return nil
		}
		data, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		var regs []Regulation
		if err := yaml.Unmarshal(data, &regs); err != nil {
			return err
		}
		for i := range regs {
			e.regulations[regs[i].Code] = &regs[i]
		}
		return nil
	})
}

func (e *Engine) loadControls(dir string) error {
	return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() || filepath.Ext(path) != ".yaml" {
			return nil
		}
		data, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		var ctrls []Control
		if err := yaml.Unmarshal(data, &ctrls); err != nil {
			return err
		}
		for i := range ctrls {
			e.controls[ctrls[i].ID] = &ctrls[i]
		}
		return nil
	})
}