Redflag/aggregator-agent/internal/crypto/verification.go

package crypto

import (
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"time"

	"github.com/Fimeg/RedFlag/aggregator-agent/internal/client"
)

// CommandVerifier handles Ed25519 signature verification for commands
type CommandVerifier struct {
	// In the future, this could include:
	// - Key rotation support
	// - Multiple trusted keys
	// - Revocation checking
}

// NewCommandVerifier creates a new command verifier
func NewCommandVerifier() *CommandVerifier {
	return &CommandVerifier{}
}

// VerifyCommand verifies that a command's signature is valid
func (v *CommandVerifier) VerifyCommand(cmd client.Command, serverPubKey ed25519.PublicKey) error {
	// Check if signature is present
	if cmd.Signature == "" {
		return fmt.Errorf("command missing signature")
	}

	// Decode the signature
	sig, err := hex.DecodeString(cmd.Signature)
	if err != nil {
		return fmt.Errorf("invalid signature encoding: %w", err)
	}

	// Verify signature length
	if len(sig) != ed25519.SignatureSize {
		return fmt.Errorf("invalid signature length: expected %d bytes, got %d",
			ed25519.SignatureSize, len(sig))
	}

	// Reconstruct the signed message
	message, err := v.reconstructMessage(cmd)
	if err != nil {
		return fmt.Errorf("failed to reconstruct message: %w", err)
	}

	// Verify the Ed25519 signature
	if !ed25519.Verify(serverPubKey, message, sig) {
		return fmt.Errorf("signature verification failed")
	}

	return nil
}

// reconstructMessage recreates the message that was signed by the server
// This must exactly match the server's signing implementation
func (v *CommandVerifier) reconstructMessage(cmd client.Command) ([]byte, error) {
	// Marshal parameters to JSON
	paramsJSON, err := json.Marshal(cmd.Params)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal parameters: %w", err)
	}

	// Create SHA256 hash of parameters
	paramsHash := sha256.Sum256(paramsJSON)
	paramsHashHex := hex.EncodeToString(paramsHash[:])

	// Create the message in the exact format the server uses
	// Format: "ID:CommandType:ParamsHash"
	message := fmt.Sprintf("%s:%s:%s",
		cmd.ID,
		cmd.Type,
		paramsHashHex)

	return []byte(message), nil
}

// VerifyCommandWithTimestamp verifies a command and checks its timestamp
// This prevents replay attacks with old commands
// Note: Timestamp verification requires the CreatedAt field which is not sent to agents
// This method is kept for future enhancement when we add timestamp to the command payload
func (v *CommandVerifier) VerifyCommandWithTimestamp(
	cmd client.Command,
	serverPubKey ed25519.PublicKey,
	maxAge time.Duration,
) error {
	// First verify the signature
	if err := v.VerifyCommand(cmd, serverPubKey); err != nil {
		return err
	}

	// Timestamp checking is currently disabled as CreatedAt is not included in the command sent to agents
	// TODO: Add CreatedAt to command payload if timestamp verification is needed

	return nil
}

// VerifyCommandBatch verifies multiple commands efficiently
// This is useful when processing multiple commands at once
func (v *CommandVerifier) VerifyCommandBatch(
	commands []client.Command,
	serverPubKey ed25519.PublicKey,
) []error {
	errors := make([]error, len(commands))

	for i, cmd := range commands {
		errors[i] = v.VerifyCommand(cmd, serverPubKey)
	}

	return errors
}

// ExtractCommandIDFromSignature attempts to verify a signature and returns the command ID
// This is useful for debugging and logging
func (v *CommandVerifier) ExtractCommandIDFromSignature(
	signature string,
	expectedMessage string,
	serverPubKey ed25519.PublicKey,
) (string, error) {
	// Decode signature
	sig, err := hex.DecodeString(signature)
	if err != nil {
		return "", fmt.Errorf("invalid signature encoding: %w", err)
	}

	// Verify signature
	if !ed25519.Verify(serverPubKey, []byte(expectedMessage), sig) {
		return "", fmt.Errorf("signature verification failed")
	}

	// In a real implementation, we might embed the command ID in the signature
	// For now, we return an empty string since the ID is part of the message
	return "", nil
}

// CheckKeyRotation checks if a public key needs to be rotated
// This is a placeholder for future key rotation support
func (v *CommandVerifier) CheckKeyRotation(currentKey ed25519.PublicKey) (ed25519.PublicKey, bool, error) {
	// In the future, this could:
	// - Check a key rotation endpoint
	// - Load multiple trusted keys
	// - Implement key pinning with fallback
	// - Handle emergency key revocation

	// For now, just return the current key
	return currentKey, false, nil
}