fetch_ml/internal/audit/audit.go

package audit

import (
	"crypto/sha256"
	"encoding/hex"
	"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"

	// HIPAA-specific file access events
	EventFileRead      EventType = "file_read"
	EventFileWrite     EventType = "file_write"
	EventFileDelete    EventType = "file_delete"
	EventDatasetAccess EventType = "dataset_access"
)

// Event represents an audit log event with integrity chain
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"` // File path, dataset ID, etc.
	Action    string                 `json:"action,omitempty"`   // read, write, delete
	Success   bool                   `json:"success"`
	ErrorMsg  string                 `json:"error,omitempty"`
	Metadata  map[string]interface{} `json:"metadata,omitempty"`

	// Integrity chain fields for tamper-evident logging (HIPAA requirement)
	PrevHash    string `json:"prev_hash,omitempty"`  // SHA-256 of previous event
	EventHash   string `json:"event_hash,omitempty"` // SHA-256 of this event
	SequenceNum int64  `json:"sequence_num,omitempty"`
}

// Logger handles audit logging with integrity chain
type Logger struct {
	enabled     bool
	filePath    string
	file        *os.File
	mu          sync.Mutex
	logger      *logging.Logger
	lastHash    string
	sequenceNum int64
}

// 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 with integrity chain
func (al *Logger) Log(event Event) {
	if !al.enabled {
		return
	}

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

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

	// Set sequence number and previous hash for integrity chain
	al.sequenceNum++
	event.SequenceNum = al.sequenceNum
	event.PrevHash = al.lastHash

	// Calculate hash of this event for tamper evidence
	event.EventHash = al.CalculateEventHash(event)
	al.lastHash = event.EventHash

	// 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,
			"seq", event.SequenceNum,
			"hash", event.EventHash[:16], // Log first 16 chars of hash
		)
	}
}

// CalculateEventHash computes SHA-256 hash of event data for integrity chain
// Exported for testing purposes
func (al *Logger) CalculateEventHash(event Event) string {
	// Create a copy without the hash field for hashing
	eventCopy := event
	eventCopy.EventHash = ""
	eventCopy.PrevHash = ""

	data, err := json.Marshal(eventCopy)
	if err != nil {
		// Fallback: hash the timestamp and type
		data = []byte(fmt.Sprintf("%s:%s:%d", event.Timestamp, event.EventType, event.SequenceNum))
	}

	hash := sha256.Sum256(data)
	return hex.EncodeToString(hash[:])
}

// LogFileAccess logs a file access operation (HIPAA requirement)
func (al *Logger) LogFileAccess(
	eventType EventType,
	userID, filePath, ipAddr string,
	success bool,
	errMsg string,
) {
	action := "read"
	switch eventType {
	case EventFileWrite:
		action = "write"
	case EventFileDelete:
		action = "delete"
	}

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

// VerifyChain checks the integrity of the audit log chain
// Returns the first sequence number where tampering is detected, or -1 if valid
func (al *Logger) VerifyChain(events []Event) (tamperedSeq int, err error) {
	if len(events) == 0 {
		return -1, nil
	}

	var expectedPrevHash string

	for _, event := range events {
		// Verify previous hash chain
		if event.SequenceNum > 1 && event.PrevHash != expectedPrevHash {
			return int(event.SequenceNum), fmt.Errorf(
				"chain break at sequence %d: expected prev_hash=%s, got %s",
				event.SequenceNum, expectedPrevHash, event.PrevHash,
			)
		}

		// Verify event hash
		expectedHash := al.CalculateEventHash(event)
		if event.EventHash != expectedHash {
			return int(event.SequenceNum), fmt.Errorf(
				"hash mismatch at sequence %d: expected %s, got %s",
				event.SequenceNum, expectedHash, event.EventHash,
			)
		}

		expectedPrevHash = event.EventHash
	}

	return -1, nil
}

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