breakpilot-pwa/ai-compliance-sdk/internal/audit/trail_builder.go

package audit

import (
	"context"
	"time"

	"github.com/google/uuid"
)

// TrailBuilder helps construct structured audit entries
type TrailBuilder struct {
	store *Store
}

// NewTrailBuilder creates a new trail builder
func NewTrailBuilder(store *Store) *TrailBuilder {
	return &TrailBuilder{store: store}
}

// LLMEntryBuilder builds LLM audit entries
type LLMEntryBuilder struct {
	entry *LLMAuditEntry
	store *Store
}

// NewLLMEntry creates a new LLM audit entry builder
func (tb *TrailBuilder) NewLLMEntry() *LLMEntryBuilder {
	return &LLMEntryBuilder{
		entry: &LLMAuditEntry{
			ID:                     uuid.New(),
			PIITypesDetected:       []string{},
			PolicyViolations:       []string{},
			DataCategoriesAccessed: []string{},
			RequestMetadata:        make(map[string]any),
			CreatedAt:              time.Now().UTC(),
		},
		store: tb.store,
	}
}

// WithTenant sets the tenant ID
func (b *LLMEntryBuilder) WithTenant(tenantID uuid.UUID) *LLMEntryBuilder {
	b.entry.TenantID = tenantID
	return b
}

// WithNamespace sets the namespace ID
func (b *LLMEntryBuilder) WithNamespace(namespaceID uuid.UUID) *LLMEntryBuilder {
	b.entry.NamespaceID = &namespaceID
	return b
}

// WithUser sets the user ID
func (b *LLMEntryBuilder) WithUser(userID uuid.UUID) *LLMEntryBuilder {
	b.entry.UserID = userID
	return b
}

// WithSession sets the session ID
func (b *LLMEntryBuilder) WithSession(sessionID string) *LLMEntryBuilder {
	b.entry.SessionID = sessionID
	return b
}

// WithOperation sets the operation type
func (b *LLMEntryBuilder) WithOperation(operation string) *LLMEntryBuilder {
	b.entry.Operation = operation
	return b
}

// WithModel sets the model used
func (b *LLMEntryBuilder) WithModel(model, provider string) *LLMEntryBuilder {
	b.entry.ModelUsed = model
	b.entry.Provider = provider
	return b
}

// WithPrompt sets prompt-related fields
func (b *LLMEntryBuilder) WithPrompt(hash string, length int) *LLMEntryBuilder {
	b.entry.PromptHash = hash
	b.entry.PromptLength = length
	return b
}

// WithResponse sets response-related fields
func (b *LLMEntryBuilder) WithResponse(length int) *LLMEntryBuilder {
	b.entry.ResponseLength = length
	return b
}

// WithUsage sets token usage and duration
func (b *LLMEntryBuilder) WithUsage(tokens int, durationMS int) *LLMEntryBuilder {
	b.entry.TokensUsed = tokens
	b.entry.DurationMS = durationMS
	return b
}

// WithPII sets PII detection fields
func (b *LLMEntryBuilder) WithPII(detected bool, types []string, redacted bool) *LLMEntryBuilder {
	b.entry.PIIDetected = detected
	b.entry.PIITypesDetected = types
	b.entry.PIIRedacted = redacted
	return b
}

// WithPolicy sets policy-related fields
func (b *LLMEntryBuilder) WithPolicy(policyID *uuid.UUID, violations []string) *LLMEntryBuilder {
	b.entry.PolicyID = policyID
	b.entry.PolicyViolations = violations
	return b
}

// WithDataCategories sets accessed data categories
func (b *LLMEntryBuilder) WithDataCategories(categories []string) *LLMEntryBuilder {
	b.entry.DataCategoriesAccessed = categories
	return b
}

// WithError sets error message
func (b *LLMEntryBuilder) WithError(errMsg string) *LLMEntryBuilder {
	b.entry.ErrorMessage = errMsg
	return b
}

// WithMetadata sets request metadata
func (b *LLMEntryBuilder) WithMetadata(metadata map[string]any) *LLMEntryBuilder {
	b.entry.RequestMetadata = metadata
	return b
}

// AddMetadata adds a key-value pair to metadata
func (b *LLMEntryBuilder) AddMetadata(key string, value any) *LLMEntryBuilder {
	if b.entry.RequestMetadata == nil {
		b.entry.RequestMetadata = make(map[string]any)
	}
	b.entry.RequestMetadata[key] = value
	return b
}

// Build returns the built entry
func (b *LLMEntryBuilder) Build() *LLMAuditEntry {
	return b.entry
}

// Save persists the entry to the database
func (b *LLMEntryBuilder) Save(ctx context.Context) error {
	return b.store.CreateLLMAuditEntry(ctx, b.entry)
}

// GeneralEntryBuilder builds general audit entries
type GeneralEntryBuilder struct {
	entry *GeneralAuditEntry
	store *Store
}

// NewGeneralEntry creates a new general audit entry builder
func (tb *TrailBuilder) NewGeneralEntry() *GeneralEntryBuilder {
	return &GeneralEntryBuilder{
		entry: &GeneralAuditEntry{
			ID:        uuid.New(),
			OldValues: make(map[string]any),
			NewValues: make(map[string]any),
			CreatedAt: time.Now().UTC(),
		},
		store: tb.store,
	}
}

// WithTenant sets the tenant ID
func (b *GeneralEntryBuilder) WithTenant(tenantID uuid.UUID) *GeneralEntryBuilder {
	b.entry.TenantID = tenantID
	return b
}

// WithNamespace sets the namespace ID
func (b *GeneralEntryBuilder) WithNamespace(namespaceID uuid.UUID) *GeneralEntryBuilder {
	b.entry.NamespaceID = &namespaceID
	return b
}

// WithUser sets the user ID
func (b *GeneralEntryBuilder) WithUser(userID uuid.UUID) *GeneralEntryBuilder {
	b.entry.UserID = userID
	return b
}

// WithAction sets the action
func (b *GeneralEntryBuilder) WithAction(action string) *GeneralEntryBuilder {
	b.entry.Action = action
	return b
}

// WithResource sets the resource type and ID
func (b *GeneralEntryBuilder) WithResource(resourceType string, resourceID *uuid.UUID) *GeneralEntryBuilder {
	b.entry.ResourceType = resourceType
	b.entry.ResourceID = resourceID
	return b
}

// WithOldValues sets the old values
func (b *GeneralEntryBuilder) WithOldValues(values map[string]any) *GeneralEntryBuilder {
	b.entry.OldValues = values
	return b
}

// WithNewValues sets the new values
func (b *GeneralEntryBuilder) WithNewValues(values map[string]any) *GeneralEntryBuilder {
	b.entry.NewValues = values
	return b
}

// WithClient sets client information
func (b *GeneralEntryBuilder) WithClient(ipAddress, userAgent string) *GeneralEntryBuilder {
	b.entry.IPAddress = ipAddress
	b.entry.UserAgent = userAgent
	return b
}

// WithReason sets the reason for the action
func (b *GeneralEntryBuilder) WithReason(reason string) *GeneralEntryBuilder {
	b.entry.Reason = reason
	return b
}

// Build returns the built entry
func (b *GeneralEntryBuilder) Build() *GeneralAuditEntry {
	return b.entry
}

// Save persists the entry to the database
func (b *GeneralEntryBuilder) Save(ctx context.Context) error {
	return b.store.CreateGeneralAuditEntry(ctx, b.entry)
}

// Common audit action types
const (
	ActionCreate    = "create"
	ActionUpdate    = "update"
	ActionDelete    = "delete"
	ActionRead      = "read"
	ActionExport    = "export"
	ActionGrant     = "grant"
	ActionRevoke    = "revoke"
	ActionLogin     = "login"
	ActionLogout    = "logout"
	ActionFailed    = "failed"
)

// Common resource types
const (
	ResourceTenant    = "tenant"
	ResourceNamespace = "namespace"
	ResourceRole      = "role"
	ResourceUserRole  = "user_role"
	ResourcePolicy    = "llm_policy"
	ResourceAPIKey    = "api_key"
	ResourceEvidence  = "evidence"
	ResourceControl   = "control"
)

// Convenience methods for common operations

// LogRoleAssignment creates an audit entry for role assignment
func (tb *TrailBuilder) LogRoleAssignment(ctx context.Context, tenantID, userID, targetUserID, roleID uuid.UUID, grantedBy uuid.UUID, ipAddress, userAgent string) error {
	return tb.NewGeneralEntry().
		WithTenant(tenantID).
		WithUser(grantedBy).
		WithAction(ActionGrant).
		WithResource(ResourceUserRole, &roleID).
		WithNewValues(map[string]any{
			"target_user_id": targetUserID.String(),
			"role_id":        roleID.String(),
		}).
		WithClient(ipAddress, userAgent).
		Save(ctx)
}

// LogRoleRevocation creates an audit entry for role revocation
func (tb *TrailBuilder) LogRoleRevocation(ctx context.Context, tenantID, userID, targetUserID, roleID uuid.UUID, revokedBy uuid.UUID, reason, ipAddress, userAgent string) error {
	return tb.NewGeneralEntry().
		WithTenant(tenantID).
		WithUser(revokedBy).
		WithAction(ActionRevoke).
		WithResource(ResourceUserRole, &roleID).
		WithOldValues(map[string]any{
			"target_user_id": targetUserID.String(),
			"role_id":        roleID.String(),
		}).
		WithReason(reason).
		WithClient(ipAddress, userAgent).
		Save(ctx)
}

// LogPolicyChange creates an audit entry for LLM policy changes
func (tb *TrailBuilder) LogPolicyChange(ctx context.Context, tenantID, userID, policyID uuid.UUID, action string, oldValues, newValues map[string]any, ipAddress, userAgent string) error {
	return tb.NewGeneralEntry().
		WithTenant(tenantID).
		WithUser(userID).
		WithAction(action).
		WithResource(ResourcePolicy, &policyID).
		WithOldValues(oldValues).
		WithNewValues(newValues).
		WithClient(ipAddress, userAgent).
		Save(ctx)
}

// LogNamespaceAccess creates an audit entry for namespace access
func (tb *TrailBuilder) LogNamespaceAccess(ctx context.Context, tenantID, userID, namespaceID uuid.UUID, action string, ipAddress, userAgent string) error {
	return tb.NewGeneralEntry().
		WithTenant(tenantID).
		WithUser(userID).
		WithNamespace(namespaceID).
		WithAction(action).
		WithResource(ResourceNamespace, &namespaceID).
		WithClient(ipAddress, userAgent).
		Save(ctx)
}

// LogDataExport creates an audit entry for data export
func (tb *TrailBuilder) LogDataExport(ctx context.Context, tenantID, userID uuid.UUID, namespaceID *uuid.UUID, resourceType, format string, recordCount int, ipAddress, userAgent string) error {
	builder := tb.NewGeneralEntry().
		WithTenant(tenantID).
		WithUser(userID).
		WithAction(ActionExport).
		WithResource(resourceType, nil).
		WithNewValues(map[string]any{
			"format":       format,
			"record_count": recordCount,
		}).
		WithClient(ipAddress, userAgent)

	if namespaceID != nil {
		builder.WithNamespace(*namespaceID)
	}

	return builder.Save(ctx)
}