fetch_ml/internal/audit/chain.go

// Package audit provides tamper-evident audit logging with hash chaining
package audit

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"sync"
)

// ChainEntry represents an audit log entry with hash chaining
type ChainEntry struct {
	Event    Event  `json:"event"`
	PrevHash string `json:"prev_hash"`
	ThisHash string `json:"this_hash"`
	SeqNum   uint64 `json:"seq_num"`
}

// HashChain maintains a chain of tamper-evident audit entries
type HashChain struct {
	mu       sync.RWMutex
	lastHash string
	seqNum   uint64
	file     *os.File
	encoder  *json.Encoder
}

// NewHashChain creates a new hash chain for audit logging
// If filePath is empty, the chain operates in-memory only
func NewHashChain(filePath string) (*HashChain, error) {
	hc := &HashChain{
		lastHash: "",
		seqNum:   0,
	}

	if filePath != "" {
		// Append to existing file if it exists, checking chain integrity
		file, seq, lastHash, err := openOrCreateChainFile(filePath)
		if err != nil {
			return nil, fmt.Errorf("failed to open chain file: %w", err)
		}
		hc.file = file
		hc.encoder = json.NewEncoder(file)
		hc.seqNum = seq
		hc.lastHash = lastHash
	}

	return hc, nil
}

// AddEvent adds an event to the hash chain
func (hc *HashChain) AddEvent(event Event) (*ChainEntry, error) {
	hc.mu.Lock()
	defer hc.mu.Unlock()

	hc.seqNum++

	entry := ChainEntry{
		Event:    event,
		PrevHash: hc.lastHash,
		SeqNum:   hc.seqNum,
	}

	// Compute hash of this entry
	data, err := json.Marshal(struct {
		Event    Event  `json:"event"`
		PrevHash string `json:"prev_hash"`
		SeqNum   uint64 `json:"seq_num"`
	}{
		Event:    entry.Event,
		PrevHash: entry.PrevHash,
		SeqNum:   entry.SeqNum,
	})
	if err != nil {
		hc.seqNum-- // Rollback sequence number
		return nil, fmt.Errorf("failed to marshal entry: %w", err)
	}

	hash := sha256.Sum256(data)
	entry.ThisHash = hex.EncodeToString(hash[:])
	hc.lastHash = entry.ThisHash

	// Write to file if configured
	if hc.encoder != nil {
		if err := hc.encoder.Encode(entry); err != nil {
			return nil, fmt.Errorf("failed to write entry: %w", err)
		}
	}

	return &entry, nil
}

// VerifyChain verifies the integrity of a chain from a file
func VerifyChain(filePath string) error {
	file, err := os.Open(filePath)
	if err != nil {
		return fmt.Errorf("failed to open chain file: %w", err)
	}
	defer file.Close()

	decoder := json.NewDecoder(file)

	var prevHash string
	var prevSeq uint64

	for lineNum := 1; decoder.More(); lineNum++ {
		var entry ChainEntry
		if err := decoder.Decode(&entry); err != nil {
			return fmt.Errorf("failed to decode entry at line %d: %w", lineNum, err)
		}

		// Verify sequence number
		if entry.SeqNum != prevSeq+1 {
			return fmt.Errorf("sequence number mismatch at entry %d: expected %d, got %d",
				lineNum, prevSeq+1, entry.SeqNum)
		}

		// Verify previous hash linkage
		if entry.PrevHash != prevHash {
			return fmt.Errorf("hash chain broken at entry %d: expected prev_hash=%s, got %s",
				lineNum, prevHash, entry.PrevHash)
		}

		// Verify this entry's hash
		data, err := json.Marshal(struct {
			Event    Event  `json:"event"`
			PrevHash string `json:"prev_hash"`
			SeqNum   uint64 `json:"seq_num"`
		}{
			Event:    entry.Event,
			PrevHash: entry.PrevHash,
			SeqNum:   entry.SeqNum,
		})
		if err != nil {
			return fmt.Errorf("failed to marshal entry %d for verification: %w", lineNum, err)
		}

		computedHash := sha256.Sum256(data)
		computedHashStr := hex.EncodeToString(computedHash[:])
		if computedHashStr != entry.ThisHash {
			return fmt.Errorf("hash mismatch at entry %d: computed=%s, stored=%s",
				lineNum, computedHashStr, entry.ThisHash)
		}

		prevHash = entry.ThisHash
		prevSeq = entry.SeqNum
	}

	return nil
}

// GetLastHash returns the current chain head hash
func (hc *HashChain) GetLastHash() string {
	hc.mu.RLock()
	defer hc.mu.RUnlock()
	return hc.lastHash
}

// GetSeqNum returns the current sequence number
func (hc *HashChain) GetSeqNum() uint64 {
	hc.mu.RLock()
	defer hc.mu.RUnlock()
	return hc.seqNum
}

// Close closes the hash chain file
func (hc *HashChain) Close() error {
	hc.mu.Lock()
	defer hc.mu.Unlock()

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

// openOrCreateChainFile opens an existing chain file and validates its integrity,
// or creates a new one if it doesn't exist
func openOrCreateChainFile(filePath string) (*os.File, uint64, string, error) {
	// Check if file exists
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		// Create new file
		file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
		if err != nil {
			return nil, 0, "", fmt.Errorf("failed to create chain file: %w", err)
		}
		return file, 0, "", nil
	}

	// File exists - verify integrity and get last state
	file, err := os.OpenFile(filePath, os.O_RDONLY, 0600)
	if err != nil {
		return nil, 0, "", fmt.Errorf("failed to open existing chain file: %w", err)
	}

	decoder := json.NewDecoder(file)
	var lastSeq uint64
	var lastHash string

	for decoder.More() {
		var entry ChainEntry
		if err := decoder.Decode(&entry); err != nil {
			file.Close()
			return nil, 0, "", fmt.Errorf("corrupted chain file: %w", err)
		}
		lastSeq = entry.SeqNum
		lastHash = entry.ThisHash
	}

	file.Close()

	// Reopen for appending
	file, err = os.OpenFile(filePath, os.O_WRONLY|os.O_APPEND, 0600)
	if err != nil {
		return nil, 0, "", fmt.Errorf("failed to reopen chain file for append: %w", err)
	}

	return file, lastSeq, lastHash, nil
}

// TamperEvidenceLogger wraps a standard audit logger with hash chaining
type TamperEvidenceLogger struct {
	logger    *Logger
	hashChain *HashChain
}

// NewTamperEvidenceLogger creates a new tamper-evident audit logger
func NewTamperEvidenceLogger(logger *Logger, chainFilePath string) (*TamperEvidenceLogger, error) {
	hashChain, err := NewHashChain(chainFilePath)
	if err != nil {
		return nil, fmt.Errorf("failed to create hash chain: %w", err)
	}

	return &TamperEvidenceLogger{
		logger:    logger,
		hashChain: hashChain,
	}, nil
}

// Log logs an event with tamper-evident hash chaining
func (tel *TamperEvidenceLogger) Log(event Event) error {
	// Log to standard logger
	if tel.logger != nil {
		tel.logger.Log(event)
	}

	// Add to hash chain
	_, err := tel.hashChain.AddEvent(event)
	if err != nil {
		return fmt.Errorf("failed to add event to hash chain: %w", err)
	}

	return nil
}

// Verify verifies the integrity of the tamper-evident log
func (tel *TamperEvidenceLogger) Verify() error {
	if tel.hashChain.file == nil {
		return fmt.Errorf("no chain file configured for verification")
	}
	return VerifyChain(tel.hashChain.file.Name())
}

// Close closes the tamper-evident logger
func (tel *TamperEvidenceLogger) Close() error {
	if tel.hashChain != nil {
		return tel.hashChain.Close()
	}
	return nil
}