skybridge/kms/internal/audit/audit.go

package audit

import (
	"context"
	"encoding/json"
	"time"

	"github.com/google/uuid"
	"go.uber.org/zap"

	"github.com/RyanCopley/skybridge/kms/internal/config"
)

// EventType represents the type of audit event
type EventType string

const (
	// Authentication events
	EventTypeLogin          EventType = "auth.login"
	EventTypeLoginFailed    EventType = "auth.login_failed"
	EventTypeLogout         EventType = "auth.logout"
	EventTypeTokenCreated   EventType = "auth.token_created"
	EventTypeTokenRevoked   EventType = "auth.token_revoked"
	EventTypeTokenValidated EventType = "auth.token_validated"
	
	// Session events
	EventTypeSessionCreated EventType = "session.created"
	EventTypeSessionRevoked EventType = "session.revoked"
	EventTypeSessionExpired EventType = "session.expired"
	
	// Application events
	EventTypeAppCreated EventType = "app.created"
	EventTypeAppUpdated EventType = "app.updated"
	EventTypeAppDeleted EventType = "app.deleted"
	
	// Permission events
	EventTypePermissionGranted EventType = "permission.granted"
	EventTypePermissionRevoked EventType = "permission.revoked"
	EventTypePermissionDenied  EventType = "permission.denied"
	
	// Tenant events
	EventTypeTenantCreated   EventType = "tenant.created"
	EventTypeTenantUpdated   EventType = "tenant.updated"
	EventTypeTenantSuspended EventType = "tenant.suspended"
	EventTypeTenantActivated EventType = "tenant.activated"
	
	// User events
	EventTypeUserCreated   EventType = "user.created"
	EventTypeUserUpdated   EventType = "user.updated"
	EventTypeUserSuspended EventType = "user.suspended"
	EventTypeUserActivated EventType = "user.activated"
	
	// Security events
	EventTypeSecurityViolation EventType = "security.violation"
	EventTypeBruteForceAttempt EventType = "security.brute_force"
	EventTypeIPBlocked         EventType = "security.ip_blocked"
	EventTypeRateLimitExceeded EventType = "security.rate_limit_exceeded"
	
	// System events
	EventTypeSystemStartup  EventType = "system.startup"
	EventTypeSystemShutdown EventType = "system.shutdown"
	EventTypeConfigChanged  EventType = "system.config_changed"
)

// EventSeverity represents the severity level of an audit event
type EventSeverity string

const (
	SeverityInfo     EventSeverity = "info"
	SeverityWarning  EventSeverity = "warning"
	SeverityError    EventSeverity = "error"
	SeverityCritical EventSeverity = "critical"
)

// EventStatus represents the status of an audit event
type EventStatus string

const (
	StatusSuccess EventStatus = "success"
	StatusFailure EventStatus = "failure"
	StatusPending EventStatus = "pending"
)

// AuditEvent represents a single audit event
type AuditEvent struct {
	ID          uuid.UUID              `json:"id" db:"id"`
	Type        EventType              `json:"type" db:"type"`
	Severity    EventSeverity          `json:"severity" db:"severity"`
	Status      EventStatus            `json:"status" db:"status"`
	Timestamp   time.Time              `json:"timestamp" db:"timestamp"`
	
	// Actor information
	ActorID     string                 `json:"actor_id,omitempty" db:"actor_id"`
	ActorType   string                 `json:"actor_type,omitempty" db:"actor_type"` // user, system, service
	ActorIP     string                 `json:"actor_ip,omitempty" db:"actor_ip"`
	UserAgent   string                 `json:"user_agent,omitempty" db:"user_agent"`
	
	// Tenant information
	TenantID    *uuid.UUID             `json:"tenant_id,omitempty" db:"tenant_id"`
	
	// Resource information
	ResourceID   string                `json:"resource_id,omitempty" db:"resource_id"`
	ResourceType string                `json:"resource_type,omitempty" db:"resource_type"`
	
	// Event details
	Action      string                 `json:"action" db:"action"`
	Description string                 `json:"description" db:"description"`
	Details     map[string]interface{} `json:"details,omitempty" db:"details"`
	
	// Request context
	RequestID   string                 `json:"request_id,omitempty" db:"request_id"`
	SessionID   string                 `json:"session_id,omitempty" db:"session_id"`
	
	// Additional metadata
	Tags        []string               `json:"tags,omitempty" db:"tags"`
	Metadata    map[string]string      `json:"metadata,omitempty" db:"metadata"`
}

// AuditLogger defines the interface for audit logging
type AuditLogger interface {
	// LogEvent logs a single audit event
	LogEvent(ctx context.Context, event *AuditEvent) error
	
	// LogAuthEvent logs an authentication-related event
	LogAuthEvent(ctx context.Context, eventType EventType, actorID, actorIP string, details map[string]interface{}) error
	
	// LogPermissionEvent logs a permission-related event
	LogPermissionEvent(ctx context.Context, eventType EventType, actorID, resourceID, resourceType string, details map[string]interface{}) error
	
	// LogSecurityEvent logs a security-related event
	LogSecurityEvent(ctx context.Context, eventType EventType, actorIP string, severity EventSeverity, details map[string]interface{}) error
	
	// LogSystemEvent logs a system-related event
	LogSystemEvent(ctx context.Context, eventType EventType, details map[string]interface{}) error
	
	// QueryEvents queries audit events with filters
	QueryEvents(ctx context.Context, filter *AuditFilter) ([]*AuditEvent, error)
	
	// GetEventByID retrieves a specific audit event by ID
	GetEventByID(ctx context.Context, eventID uuid.UUID) (*AuditEvent, error)
	
	// GetEventStats returns audit event statistics
	GetEventStats(ctx context.Context, filter *AuditStatsFilter) (*AuditStats, error)
}

// AuditFilter represents filters for querying audit events
type AuditFilter struct {
	EventTypes   []EventType    `json:"event_types,omitempty"`
	Severities   []EventSeverity `json:"severities,omitempty"`
	Statuses     []EventStatus  `json:"statuses,omitempty"`
	ActorID      string         `json:"actor_id,omitempty"`
	ActorType    string         `json:"actor_type,omitempty"`
	TenantID     *uuid.UUID     `json:"tenant_id,omitempty"`
	ResourceID   string         `json:"resource_id,omitempty"`
	ResourceType string         `json:"resource_type,omitempty"`
	StartTime    *time.Time     `json:"start_time,omitempty"`
	EndTime      *time.Time     `json:"end_time,omitempty"`
	Tags         []string       `json:"tags,omitempty"`
	Limit        int            `json:"limit"`
	Offset       int            `json:"offset"`
	OrderBy      string         `json:"order_by"` // timestamp, type, severity
	OrderDesc    bool           `json:"order_desc"`
}

// AuditStatsFilter represents filters for audit statistics
type AuditStatsFilter struct {
	EventTypes []EventType `json:"event_types,omitempty"`
	TenantID   *uuid.UUID  `json:"tenant_id,omitempty"`
	StartTime  *time.Time  `json:"start_time,omitempty"`
	EndTime    *time.Time  `json:"end_time,omitempty"`
	GroupBy    string      `json:"group_by"` // type, severity, status, hour, day
}

// AuditStats represents audit event statistics
type AuditStats struct {
	TotalEvents int                    `json:"total_events"`
	ByType      map[EventType]int      `json:"by_type"`
	BySeverity  map[EventSeverity]int  `json:"by_severity"`
	ByStatus    map[EventStatus]int    `json:"by_status"`
	ByTime      map[string]int         `json:"by_time,omitempty"`
}

// auditLogger implements the AuditLogger interface
type auditLogger struct {
	config     config.ConfigProvider
	logger     *zap.Logger
	repository AuditRepository
}

// AuditRepository defines the interface for audit event storage
type AuditRepository interface {
	Create(ctx context.Context, event *AuditEvent) error
	Query(ctx context.Context, filter *AuditFilter) ([]*AuditEvent, error)
	GetByID(ctx context.Context, eventID uuid.UUID) (*AuditEvent, error)
	GetStats(ctx context.Context, filter *AuditStatsFilter) (*AuditStats, error)
	DeleteOldEvents(ctx context.Context, olderThan time.Time) (int, error)
}

// NewAuditLogger creates a new audit logger
func NewAuditLogger(config config.ConfigProvider, logger *zap.Logger, repository AuditRepository) AuditLogger {
	return &auditLogger{
		config:     config,
		logger:     logger,
		repository: repository,
	}
}

// LogEvent logs a single audit event
func (a *auditLogger) LogEvent(ctx context.Context, event *AuditEvent) error {
	// Set default values
	if event.ID == uuid.Nil {
		event.ID = uuid.New()
	}
	if event.Timestamp.IsZero() {
		event.Timestamp = time.Now().UTC()
	}
	if event.Severity == "" {
		event.Severity = SeverityInfo
	}
	if event.Status == "" {
		event.Status = StatusSuccess
	}

	// Extract request context if available
	if requestID := ctx.Value("request_id"); requestID != nil {
		if reqID, ok := requestID.(string); ok {
			event.RequestID = reqID
		}
	}

	// Log to structured logger
	a.logToStructuredLogger(event)

	// Store in repository
	if err := a.repository.Create(ctx, event); err != nil {
		a.logger.Error("Failed to store audit event", 
			zap.Error(err),
			zap.String("event_id", event.ID.String()),
			zap.String("event_type", string(event.Type)))
		return err
	}

	return nil
}

// LogAuthEvent logs an authentication-related event
func (a *auditLogger) LogAuthEvent(ctx context.Context, eventType EventType, actorID, actorIP string, details map[string]interface{}) error {
	severity := SeverityInfo
	status := StatusSuccess
	
	// Determine severity and status based on event type
	switch eventType {
	case EventTypeLoginFailed:
		severity = SeverityWarning
		status = StatusFailure
	case EventTypeTokenRevoked:
		severity = SeverityWarning
	}

	event := &AuditEvent{
		Type:        eventType,
		Severity:    severity,
		Status:      status,
		ActorID:     actorID,
		ActorType:   "user",
		ActorIP:     actorIP,
		Action:      string(eventType),
		Description: a.generateDescription(eventType, details),
		Details:     details,
		Tags:        []string{"authentication"},
	}

	return a.LogEvent(ctx, event)
}

// LogPermissionEvent logs a permission-related event
func (a *auditLogger) LogPermissionEvent(ctx context.Context, eventType EventType, actorID, resourceID, resourceType string, details map[string]interface{}) error {
	severity := SeverityInfo
	status := StatusSuccess
	
	// Determine severity and status based on event type
	switch eventType {
	case EventTypePermissionDenied:
		severity = SeverityWarning
		status = StatusFailure
	}

	event := &AuditEvent{
		Type:         eventType,
		Severity:     severity,
		Status:       status,
		ActorID:      actorID,
		ActorType:    "user",
		ResourceID:   resourceID,
		ResourceType: resourceType,
		Action:       string(eventType),
		Description:  a.generateDescription(eventType, details),
		Details:      details,
		Tags:         []string{"permission", "authorization"},
	}

	return a.LogEvent(ctx, event)
}

// LogSecurityEvent logs a security-related event
func (a *auditLogger) LogSecurityEvent(ctx context.Context, eventType EventType, actorIP string, severity EventSeverity, details map[string]interface{}) error {
	status := StatusSuccess
	if severity == SeverityError || severity == SeverityCritical {
		status = StatusFailure
	}

	event := &AuditEvent{
		Type:        eventType,
		Severity:    severity,
		Status:      status,
		ActorIP:     actorIP,
		ActorType:   "system",
		Action:      string(eventType),
		Description: a.generateDescription(eventType, details),
		Details:     details,
		Tags:        []string{"security"},
	}

	return a.LogEvent(ctx, event)
}

// LogSystemEvent logs a system-related event
func (a *auditLogger) LogSystemEvent(ctx context.Context, eventType EventType, details map[string]interface{}) error {
	event := &AuditEvent{
		Type:        eventType,
		Severity:    SeverityInfo,
		Status:      StatusSuccess,
		ActorType:   "system",
		Action:      string(eventType),
		Description: a.generateDescription(eventType, details),
		Details:     details,
		Tags:        []string{"system"},
	}

	return a.LogEvent(ctx, event)
}

// QueryEvents queries audit events with filters
func (a *auditLogger) QueryEvents(ctx context.Context, filter *AuditFilter) ([]*AuditEvent, error) {
	// Set default pagination
	if filter.Limit <= 0 {
		filter.Limit = 100
	}
	if filter.Limit > 1000 {
		filter.Limit = 1000
	}
	if filter.OrderBy == "" {
		filter.OrderBy = "timestamp"
		filter.OrderDesc = true
	}

	return a.repository.Query(ctx, filter)
}

// GetEventByID retrieves a specific audit event by ID
func (a *auditLogger) GetEventByID(ctx context.Context, eventID uuid.UUID) (*AuditEvent, error) {
	return a.repository.GetByID(ctx, eventID)
}

// GetEventStats returns audit event statistics
func (a *auditLogger) GetEventStats(ctx context.Context, filter *AuditStatsFilter) (*AuditStats, error) {
	return a.repository.GetStats(ctx, filter)
}

// logToStructuredLogger logs the event to the structured logger
func (a *auditLogger) logToStructuredLogger(event *AuditEvent) {
	fields := []zap.Field{
		zap.String("audit_event_id", event.ID.String()),
		zap.String("event_type", string(event.Type)),
		zap.String("severity", string(event.Severity)),
		zap.String("status", string(event.Status)),
		zap.Time("timestamp", event.Timestamp),
		zap.String("action", event.Action),
		zap.String("description", event.Description),
	}

	if event.ActorID != "" {
		fields = append(fields, zap.String("actor_id", event.ActorID))
	}
	if event.ActorType != "" {
		fields = append(fields, zap.String("actor_type", event.ActorType))
	}
	if event.ActorIP != "" {
		fields = append(fields, zap.String("actor_ip", event.ActorIP))
	}
	if event.TenantID != nil {
		fields = append(fields, zap.String("tenant_id", event.TenantID.String()))
	}
	if event.ResourceID != "" {
		fields = append(fields, zap.String("resource_id", event.ResourceID))
	}
	if event.ResourceType != "" {
		fields = append(fields, zap.String("resource_type", event.ResourceType))
	}
	if event.RequestID != "" {
		fields = append(fields, zap.String("request_id", event.RequestID))
	}
	if event.SessionID != "" {
		fields = append(fields, zap.String("session_id", event.SessionID))
	}
	if len(event.Tags) > 0 {
		fields = append(fields, zap.Strings("tags", event.Tags))
	}
	if event.Details != nil {
		if detailsJSON, err := json.Marshal(event.Details); err == nil {
			fields = append(fields, zap.String("details", string(detailsJSON)))
		}
	}

	// Log at appropriate level based on severity
	switch event.Severity {
	case SeverityInfo:
		a.logger.Info("Audit event", fields...)
	case SeverityWarning:
		a.logger.Warn("Audit event", fields...)
	case SeverityError:
		a.logger.Error("Audit event", fields...)
	case SeverityCritical:
		a.logger.Error("Critical audit event", fields...)
	default:
		a.logger.Info("Audit event", fields...)
	}
}

// generateDescription generates a human-readable description for an event
func (a *auditLogger) generateDescription(eventType EventType, details map[string]interface{}) string {
	switch eventType {
	case EventTypeLogin:
		return "User successfully logged in"
	case EventTypeLoginFailed:
		return "User login attempt failed"
	case EventTypeLogout:
		return "User logged out"
	case EventTypeTokenCreated:
		return "API token created"
	case EventTypeTokenRevoked:
		return "API token revoked"
	case EventTypeTokenValidated:
		return "API token validated"
	case EventTypeSessionCreated:
		return "User session created"
	case EventTypeSessionRevoked:
		return "User session revoked"
	case EventTypeSessionExpired:
		return "User session expired"
	case EventTypeAppCreated:
		return "Application created"
	case EventTypeAppUpdated:
		return "Application updated"
	case EventTypeAppDeleted:
		return "Application deleted"
	case EventTypePermissionGranted:
		return "Permission granted"
	case EventTypePermissionRevoked:
		return "Permission revoked"
	case EventTypePermissionDenied:
		return "Permission denied"
	case EventTypeTenantCreated:
		return "Tenant created"
	case EventTypeTenantUpdated:
		return "Tenant updated"
	case EventTypeTenantSuspended:
		return "Tenant suspended"
	case EventTypeTenantActivated:
		return "Tenant activated"
	case EventTypeUserCreated:
		return "User created"
	case EventTypeUserUpdated:
		return "User updated"
	case EventTypeUserSuspended:
		return "User suspended"
	case EventTypeUserActivated:
		return "User activated"
	case EventTypeSecurityViolation:
		return "Security violation detected"
	case EventTypeBruteForceAttempt:
		return "Brute force attempt detected"
	case EventTypeIPBlocked:
		return "IP address blocked"
	case EventTypeRateLimitExceeded:
		return "Rate limit exceeded"
	case EventTypeSystemStartup:
		return "System started"
	case EventTypeSystemShutdown:
		return "System shutdown"
	case EventTypeConfigChanged:
		return "Configuration changed"
	default:
		return string(eventType)
	}
}

// AuditEventBuilder provides a fluent interface for building audit events
type AuditEventBuilder struct {
	event *AuditEvent
}

// NewAuditEventBuilder creates a new audit event builder
func NewAuditEventBuilder(eventType EventType) *AuditEventBuilder {
	return &AuditEventBuilder{
		event: &AuditEvent{
			ID:        uuid.New(),
			Type:      eventType,
			Timestamp: time.Now().UTC(),
			Severity:  SeverityInfo,
			Status:    StatusSuccess,
			Details:   make(map[string]interface{}),
			Metadata:  make(map[string]string),
		},
	}
}

// WithSeverity sets the event severity
func (b *AuditEventBuilder) WithSeverity(severity EventSeverity) *AuditEventBuilder {
	b.event.Severity = severity
	return b
}

// WithStatus sets the event status
func (b *AuditEventBuilder) WithStatus(status EventStatus) *AuditEventBuilder {
	b.event.Status = status
	return b
}

// WithActor sets the actor information
func (b *AuditEventBuilder) WithActor(actorID, actorType, actorIP string) *AuditEventBuilder {
	b.event.ActorID = actorID
	b.event.ActorType = actorType
	b.event.ActorIP = actorIP
	return b
}

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

// WithResource sets the resource information
func (b *AuditEventBuilder) WithResource(resourceID, resourceType string) *AuditEventBuilder {
	b.event.ResourceID = resourceID
	b.event.ResourceType = resourceType
	return b
}

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

// WithDescription sets the description
func (b *AuditEventBuilder) WithDescription(description string) *AuditEventBuilder {
	b.event.Description = description
	return b
}

// WithDetail adds a detail
func (b *AuditEventBuilder) WithDetail(key string, value interface{}) *AuditEventBuilder {
	b.event.Details[key] = value
	return b
}

// WithDetails sets multiple details
func (b *AuditEventBuilder) WithDetails(details map[string]interface{}) *AuditEventBuilder {
	for k, v := range details {
		b.event.Details[k] = v
	}
	return b
}

// WithRequestContext sets request context information
func (b *AuditEventBuilder) WithRequestContext(requestID, sessionID string) *AuditEventBuilder {
	b.event.RequestID = requestID
	b.event.SessionID = sessionID
	return b
}

// WithTags sets the tags
func (b *AuditEventBuilder) WithTags(tags ...string) *AuditEventBuilder {
	b.event.Tags = tags
	return b
}

// WithMetadata adds metadata
func (b *AuditEventBuilder) WithMetadata(key, value string) *AuditEventBuilder {
	b.event.Metadata[key] = value
	return b
}

// Build returns the built audit event
func (b *AuditEventBuilder) Build() *AuditEvent {
	return b.event
}