fetch_ml/internal/audit/audit.go

package audit

import (
	"encoding/json"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/jfraeys/fetch_ml/internal/logging"
)

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

const (
	EventAuthAttempt       EventType = "authentication_attempt"
	EventAuthSuccess       EventType = "authentication_success"
	EventAuthFailure       EventType = "authentication_failure"
	EventJobQueued         EventType = "job_queued"
	EventJobStarted        EventType = "job_started"
	EventJobCompleted      EventType = "job_completed"
	EventJobFailed         EventType = "job_failed"
	EventJupyterStart      EventType = "jupyter_start"
	EventJupyterStop       EventType = "jupyter_stop"
	EventExperimentCreated EventType = "experiment_created"
	EventExperimentDeleted EventType = "experiment_deleted"
)

// Event represents an audit log event
type Event struct {
	Timestamp time.Time              `json:"timestamp"`
	EventType EventType              `json:"event_type"`
	UserID    string                 `json:"user_id,omitempty"`
	IPAddress string                 `json:"ip_address,omitempty"`
	Resource  string                 `json:"resource,omitempty"`
	Action    string                 `json:"action,omitempty"`
	Success   bool                   `json:"success"`
	ErrorMsg  string                 `json:"error,omitempty"`
	Metadata  map[string]interface{} `json:"metadata,omitempty"`
}

// Logger handles audit logging
type Logger struct {
	enabled  bool
	filePath string
	file     *os.File
	mu       sync.Mutex
	logger   *logging.Logger
}

// NewLogger creates a new audit logger
func NewLogger(enabled bool, filePath string, logger *logging.Logger) (*Logger, error) {
	al := &Logger{
		enabled:  enabled,
		filePath: filePath,
		logger:   logger,
	}

	if enabled && filePath != "" {
		file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
		if err != nil {
			return nil, fmt.Errorf("failed to open audit log file: %w", err)
		}
		al.file = file
	}

	return al, nil
}

// Log logs an audit event
func (al *Logger) Log(event Event) {
	if !al.enabled {
		return
	}

	event.Timestamp = time.Now().UTC()

	al.mu.Lock()
	defer al.mu.Unlock()

	// Marshal to JSON
	data, err := json.Marshal(event)
	if err != nil {
		if al.logger != nil {
			al.logger.Error("failed to marshal audit event", "error", err)
		}
		return
	}

	// Write to file if configured
	if al.file != nil {
		_, err = al.file.Write(append(data, '\n'))
		if err != nil && al.logger != nil {
			al.logger.Error("failed to write audit event", "error", err)
		}
	}

	// Also log via structured logger
	if al.logger != nil {
		al.logger.Info("audit_event",
			"event_type", event.EventType,
			"user_id", event.UserID,
			"resource", event.Resource,
			"success", event.Success,
		)
	}
}

// LogAuthAttempt logs an authentication attempt
func (al *Logger) LogAuthAttempt(userID, ipAddr string, success bool, errMsg string) {
	eventType := EventAuthSuccess
	if !success {
		eventType = EventAuthFailure
	}

	al.Log(Event{
		EventType: eventType,
		UserID:    userID,
		IPAddress: ipAddr,
		Success:   success,
		ErrorMsg:  errMsg,
	})
}

// LogJobOperation logs a job-related operation
func (al *Logger) LogJobOperation(
	eventType EventType,
	userID, jobID, ipAddr string,
	success bool,
	errMsg string,
) {
	al.Log(Event{
		EventType: eventType,
		UserID:    userID,
		IPAddress: ipAddr,
		Resource:  jobID,
		Action:    "job_operation",
		Success:   success,
		ErrorMsg:  errMsg,
	})
}

// LogJupyterOperation logs a Jupyter service operation
func (al *Logger) LogJupyterOperation(
	eventType EventType,
	userID, serviceID, ipAddr string,
	success bool,
	errMsg string,
) {
	al.Log(Event{
		EventType: eventType,
		UserID:    userID,
		IPAddress: ipAddr,
		Resource:  serviceID,
		Action:    "jupyter_operation",
		Success:   success,
		ErrorMsg:  errMsg,
	})
}

// Close closes the audit logger
func (al *Logger) Close() error {
	al.mu.Lock()
	defer al.mu.Unlock()

	if al.file != nil {
		return al.file.Close()
	}
	return nil
}