mpc/pkg/identity/identity.go

package identity

import (
	"crypto/ed25519"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"sync"
	"syscall"

	"filippo.io/age"
	"golang.org/x/term"

	"github.com/spf13/viper"

	"github.com/luxfi/mpc/pkg/common/pathutil"
	"github.com/luxfi/mpc/pkg/logger"
	"github.com/luxfi/mpc/pkg/types"
)

// NodeIdentity represents a node's identity information
type NodeIdentity struct {
	NodeName  string `json:"node_name"`
	NodeID    string `json:"node_id"`
	PublicKey string `json:"public_key"`
	CreatedAt string `json:"created_at"`
}

// Store manages node identities
type Store interface {
	// GetPublicKey retrieves a node's public key by its ID
	GetPublicKey(nodeID string) ([]byte, error)
	VerifyInitiatorMessage(msg types.InitiatorMessage) error
	// SignWireMessage signs a protocol wire message with this node's Ed25519 key.
	SignWireMessage(msg *types.Message)
	// VerifyWireMessage verifies a protocol wire message's Ed25519 signature
	// using the sender's public key looked up by SenderNodeID.
	VerifyWireMessage(msg *types.Message) error
}

// fileStore implements the Store interface using the filesystem
type fileStore struct {
	identityDir     string
	currentNodeName string

	// Cache for public keys by node_id
	publicKeys map[string][]byte
	mu         sync.RWMutex

	// Cached private key
	privateKey      []byte
	initiatorPubKey []byte
}

// NewFileStore creates a new identity store
func NewFileStore(identityDir, nodeName string, decrypt bool) (*fileStore, error) {
	if err := os.MkdirAll(identityDir, 0750); err != nil {
		return nil, fmt.Errorf("failed to create identity directory: %w", err)
	}

	privateKeyHex, err := loadPrivateKey(identityDir, nodeName, decrypt)
	if err != nil {
		return nil, err
	}

	privateKey, err := hex.DecodeString(privateKeyHex)
	if err != nil {
		return nil, fmt.Errorf("invalid private key format: %w", err)
	}

	pubKeyHex := viper.GetString("event_initiator_pubkey")
	if pubKeyHex == "" {
		return nil, fmt.Errorf("event_initiator_pubkey not found in quax config")
	}
	initiatorPubKey, err := hex.DecodeString(pubKeyHex)
	if err != nil {
		return nil, fmt.Errorf("invalid initiator public key format: %w", err)
	}

	logger.Infof("Loaded initiator public key for node %s", pubKeyHex)

	// Load peers.json to validate all nodes have identity files
	peersData, err := os.ReadFile("peers.json")
	if err != nil {
		return nil, fmt.Errorf("failed to read peers.json: %w", err)
	}

	peers := make(map[string]string)
	if err := json.Unmarshal(peersData, &peers); err != nil {
		return nil, fmt.Errorf("failed to parse peers.json: %w", err)
	}

	store := &fileStore{
		identityDir:     identityDir,
		currentNodeName: nodeName,
		publicKeys:      make(map[string][]byte),
		privateKey:      privateKey,
		initiatorPubKey: initiatorPubKey,
	}

	// Check that each node in peers.json has an identity file
	for nodeName, nodeID := range peers {
		identityFileName := fmt.Sprintf("%s_identity.json", nodeName)
		identityFilePath, err := pathutil.SafePath(identityDir, identityFileName)
		if err != nil {
			return nil, fmt.Errorf("invalid identity file path for node %s: %w", nodeName, err)
		}

		data, err := os.ReadFile(identityFilePath)
		if err != nil {
			return nil, fmt.Errorf("missing identity file for node %s (%s): %w", nodeName, nodeID, err)
		}

		var identity NodeIdentity
		if err := json.Unmarshal(data, &identity); err != nil {
			return nil, fmt.Errorf("failed to parse identity file for node %s: %w", nodeName, err)
		}

		// Verify that the nodeID in peers.json matches the one in the identity file
		if identity.NodeID != nodeID {
			return nil, fmt.Errorf("node ID mismatch for %s: %s in peers.json vs %s in identity file",
				nodeName, nodeID, identity.NodeID)
		}

		key, err := hex.DecodeString(identity.PublicKey)
		if err != nil {
			return nil, fmt.Errorf("invalid public key format for node %s: %w", nodeName, err)
		}

		store.publicKeys[identity.NodeID] = key
	}

	return store, nil
}

// loadPrivateKey loads the private key from file, decrypting if necessary
func loadPrivateKey(identityDir, nodeName string, decrypt bool) (string, error) {
	// Check for encrypted or unencrypted private key
	encryptedKeyFileName := fmt.Sprintf("%s_private.key.age", nodeName)
	unencryptedKeyFileName := fmt.Sprintf("%s_private.key", nodeName)

	encryptedKeyPath, err := pathutil.SafePath(identityDir, encryptedKeyFileName)
	if err != nil {
		return "", fmt.Errorf("invalid encrypted key path for node %s: %w", nodeName, err)
	}

	unencryptedKeyPath, err := pathutil.SafePath(identityDir, unencryptedKeyFileName)
	if err != nil {
		return "", fmt.Errorf("invalid unencrypted key path for node %s: %w", nodeName, err)
	}

	if decrypt {
		// Use the encrypted age file
		if _, err := os.Stat(encryptedKeyPath); err != nil {
			return "", fmt.Errorf("no encrypted private key found for node %s", nodeName)
		}

		logger.Infof("Using age-encrypted private key for %s", nodeName)

		// Open the encrypted file
		encryptedFile, err := os.Open(encryptedKeyPath)
		if err != nil {
			return "", fmt.Errorf("failed to open encrypted key file: %w", err)
		}
		defer encryptedFile.Close()

		// Prompt for passphrase using term.ReadPassword
		fmt.Print("Enter passphrase to decrypt private key: ")
		bytePassword, err := term.ReadPassword(syscall.Stdin)
		fmt.Println() // newline after prompt
		if err != nil {
			return "", fmt.Errorf("failed to read passphrase: %w", err)
		}
		passphrase := string(bytePassword)
		// Create an identity with the provided passphrase
		identity, err := age.NewScryptIdentity(passphrase)
		if err != nil {
			return "", fmt.Errorf("failed to create identity for decryption: %w", err)
		}

		// Decrypt the file
		decrypter, err := age.Decrypt(encryptedFile, identity)
		if err != nil {
			return "", fmt.Errorf("failed to decrypt private key: %w", err)
		}

		// Read the decrypted content
		decryptedData, err := io.ReadAll(decrypter)
		if err != nil {
			return "", fmt.Errorf("failed to read decrypted key: %w", err)
		}

		return string(decryptedData), nil
	} else {
		// Use the unencrypted private key file
		if _, err := os.Stat(unencryptedKeyPath); err != nil {
			return "", fmt.Errorf("no unencrypted private key found for node %s", nodeName)
		}

		logger.Infof("Using unencrypted private key for %s", nodeName)
		privateKeyData, err := os.ReadFile(unencryptedKeyPath)
		if err != nil {
			return "", fmt.Errorf("failed to read private key file: %w", err)
		}
		return string(privateKeyData), nil
	}
}

// GetPublicKey retrieves a node's public key by its ID
func (s *fileStore) GetPublicKey(nodeID string) ([]byte, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if key, exists := s.publicKeys[nodeID]; exists {
		return key, nil
	}

	return nil, fmt.Errorf("public key not found for node ID: %s", nodeID)
}

// SignWireMessage signs a wire message with this node's Ed25519 private key.
func (s *fileStore) SignWireMessage(msg *types.Message) {
	msg.Sign(ed25519.PrivateKey(s.privateKey))
}

// VerifyWireMessage verifies a wire message's Ed25519 signature using the
// sender's public key looked up by SenderNodeID.
func (s *fileStore) VerifyWireMessage(msg *types.Message) error {
	if len(msg.Signature) == 0 {
		return fmt.Errorf("message has no signature")
	}
	nodeID := msg.SenderNodeID
	if nodeID == "" {
		// Fall back to extracting from SenderID (party ID format "nodeID:purpose:version")
		nodeID = msg.SenderID
	}
	publicKey, err := s.GetPublicKey(nodeID)
	if err != nil {
		return fmt.Errorf("failed to get sender's public key for node %s: %w", nodeID, err)
	}
	return msg.Verify(ed25519.PublicKey(publicKey))
}

// VerifyInitiatorMessage verifies that a message was signed by the known initiator
func (s *fileStore) VerifyInitiatorMessage(msg types.InitiatorMessage) error {
	// Get the raw message that was signed
	msgBytes, err := msg.Raw()
	if err != nil {
		return fmt.Errorf("failed to get raw message data: %w", err)
	}

	// Get the signature
	signature := msg.Sig()
	if len(signature) == 0 {
		return errors.New("signature is empty")
	}

	// Verify the signature using the initiator's public key
	if !ed25519.Verify(s.initiatorPubKey, msgBytes, signature) {
		return fmt.Errorf("invalid signature from initiator")
	}

	return nil
}