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) }