pop/internal/mail/pgp.go

package mail

import (
	"crypto/rand"
	"fmt"
	"sync"

	"github.com/ProtonMail/gopenpgp/v2/crypto"
)

type PGPKeyRing struct {
	mu             sync.Mutex
	PrivateKey     *crypto.Key
	PublicKey      []byte
	PrivateKeyData []byte
}

type PGPService struct {
	keyRing *PGPKeyRing
}

func NewPGPService(privateKeyArmored string) (*PGPService, error) {
	privateKey, err := crypto.NewKeyFromArmored(privateKeyArmored)
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}

	publicKey, err := privateKey.GetPublicKey()
	if err != nil {
		return nil, fmt.Errorf("failed to extract public key: %w", err)
	}

	return &PGPService{
		keyRing: &PGPKeyRing{
			PrivateKey:     privateKey,
			PublicKey:      publicKey,
			PrivateKeyData: []byte(privateKeyArmored),
		},
	}, nil
}

func (s *PGPService) Encrypt(plaintext string, recipientPublicKey *crypto.Key) (string, error) {
	pgpMessage := crypto.NewPlainMessage([]byte(plaintext))

	recipientKeyRing, err := crypto.NewKeyRing(recipientPublicKey)
	if err != nil {
		return "", fmt.Errorf("failed to create recipient key ring: %w", err)
	}

	encrypted, err := recipientKeyRing.Encrypt(pgpMessage, nil)
	if err != nil {
		return "", fmt.Errorf("failed to encrypt: %w", err)
	}

	armored, err := encrypted.GetArmored()
	if err != nil {
		return "", fmt.Errorf("failed to armor encrypted message: %w", err)
	}

	return armored, nil
}

func (s *PGPService) EncryptBody(plaintext string, passphrase string) (string, error) {
	pgpMessage := crypto.NewPlainMessage([]byte(plaintext))

	pubKeyBytes, err := s.keyRing.PrivateKey.GetPublicKey()
	if err != nil {
		return "", fmt.Errorf("failed to get public key: %w", err)
	}

	pubKey, err := crypto.NewKeyFromArmored(string(pubKeyBytes))
	if err != nil {
		return "", fmt.Errorf("failed to parse public key: %w", err)
	}

	recipientKeyRing, err := crypto.NewKeyRing(pubKey)
	if err != nil {
		return "", fmt.Errorf("failed to create encryption key ring: %w", err)
	}

	signingKeyRing, err := s.getUnlockedKeyRing(passphrase)
	if err != nil {
		return "", fmt.Errorf("failed to create signing key ring: %w", err)
	}

	encrypted, err := recipientKeyRing.Encrypt(pgpMessage, signingKeyRing)
	if err != nil {
		return "", fmt.Errorf("failed to encrypt body: %w", err)
	}

	armored, err := encrypted.GetArmored()
	if err != nil {
		return "", fmt.Errorf("failed to armor encrypted body: %w", err)
	}

	return armored, nil
}

func (s *PGPService) EncryptAndSign(plaintext string, recipientPublicKey *crypto.Key, passphrase string) (string, error) {
	pgpMessage := crypto.NewPlainMessage([]byte(plaintext))

	recipientKeyRing, err := crypto.NewKeyRing(recipientPublicKey)
	if err != nil {
		return "", fmt.Errorf("failed to create recipient key ring: %w", err)
	}

	signingKeyRing, err := s.getUnlockedKeyRing(passphrase)
	if err != nil {
		return "", fmt.Errorf("failed to create signing key ring: %w", err)
	}

	encrypted, err := recipientKeyRing.Encrypt(pgpMessage, signingKeyRing)
	if err != nil {
		return "", fmt.Errorf("failed to encrypt and sign: %w", err)
	}

	armored, err := encrypted.GetArmored()
	if err != nil {
		return "", fmt.Errorf("failed to armor encrypted message: %w", err)
	}

	return armored, nil
}

func (s *PGPService) getUnlockedKeyRing(passphrase string) (*crypto.KeyRing, error) {
	s.keyRing.mu.Lock()
	key, err := crypto.NewKeyFromArmored(string(s.keyRing.PrivateKeyData))
	s.keyRing.mu.Unlock()
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}

	if passphrase != "" {
		unlockedKey, err := key.Unlock([]byte(passphrase))
		if err != nil {
			return nil, fmt.Errorf("failed to unlock private key: %w", err)
		}
		key = unlockedKey
	}

	return crypto.NewKeyRing(key)
}

func (s *PGPService) Decrypt(encrypted string, passphrase string) (string, error) {
	pgpMessage, err := crypto.NewPGPMessageFromArmored(encrypted)
	if err != nil {
		return "", fmt.Errorf("failed to parse encrypted message: %w", err)
	}

	decryptionKeyRing, err := s.getUnlockedKeyRing(passphrase)
	if err != nil {
		return "", fmt.Errorf("failed to create decryption key ring: %w", err)
	}

	decrypted, err := decryptionKeyRing.Decrypt(pgpMessage, nil, 0)
	if err != nil {
		return "", fmt.Errorf("failed to decrypt: %w", err)
	}

	return decrypted.GetString(), nil
}

func (s *PGPService) GenerateKeyPair(email string, passphrase string) (privateKey, publicKey string, err error) {
	key, err := crypto.GenerateKey(email, passphrase, "RSA", 4096)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate key pair: %w", err)
	}

	privateArmor, err := key.Armor()
	if err != nil {
		return "", "", fmt.Errorf("failed to armor private key: %w", err)
	}

	pubKeyBytes, err := key.GetPublicKey()
	if err != nil {
		return "", "", fmt.Errorf("failed to extract public key: %w", err)
	}

	pubArmor := string(pubKeyBytes)

	return privateArmor, pubArmor, nil
}

func (s *PGPService) GetFingerprint() (string, error) {
	if s.keyRing == nil || s.keyRing.PrivateKey == nil {
		return "", fmt.Errorf("no key ring available")
	}
	fingerprint := s.keyRing.PrivateKey.GetFingerprint()
	return fingerprint, nil
}

func (s *PGPService) ZeroPrivateKeyData() {
	if s.keyRing == nil {
		return
	}
	s.keyRing.mu.Lock()
	defer s.keyRing.mu.Unlock()
	for i := range s.keyRing.PrivateKeyData {
		s.keyRing.PrivateKeyData[i] = 0
	}
}

func (s *PGPService) SignData(data []byte, passphrase string) (string, error) {
	pgpMessage := crypto.NewPlainMessage(data)

	signingKeyRing, err := s.getUnlockedKeyRing(passphrase)
	if err != nil {
		return "", fmt.Errorf("failed to create signing key ring: %w", err)
	}

	signed, err := signingKeyRing.SignDetached(pgpMessage)
	if err != nil {
		return "", fmt.Errorf("failed to sign data: %w", err)
	}

	armored, err := signed.GetArmored()
	if err != nil {
		return "", fmt.Errorf("failed to armor signed data: %w", err)
	}

	return armored, nil
}

func (s *PGPService) EncryptAttachment(data []byte, recipientPublicKey *crypto.Key) (*Attachment, error) {
	symKey := make([]byte, 32)
	if _, err := rand.Read(symKey); err != nil {
		return nil, fmt.Errorf("failed to generate symmetric key: %w", err)
	}

	pgpMessage := crypto.NewPlainMessage(data)

	sk, err := crypto.NewSessionKeyFromToken(symKey, "AES256").Encrypt(pgpMessage)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt attachment: %w", err)
	}

	// Encrypt the symmetric key with recipient's public key
	recipientKeyRing, err := crypto.NewKeyRing(recipientPublicKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create recipient key ring: %w", err)
	}

	encryptedSymKey, err := recipientKeyRing.EncryptSessionKey(
		crypto.NewSessionKeyFromToken(symKey, "AES256"),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt symmetric key: %w", err)
	}

	return &Attachment{
		DataEnc: string(sk),
		Keys: []AttachmentKey{{
			DataEnc: string(encryptedSymKey),
		}},
	}, nil
}

func (s *PGPService) DecryptAttachment(attachment *Attachment, passphrase string) ([]byte, error) {
	if len(attachment.Keys) == 0 {
		return nil, fmt.Errorf("no keys available for attachment decryption")
	}

	decryptionKeyRing, err := s.getUnlockedKeyRing(passphrase)
	if err != nil {
		return nil, fmt.Errorf("failed to create decryption key ring: %w", err)
	}

	sk, err := decryptionKeyRing.DecryptSessionKey([]byte(attachment.Keys[0].DataEnc))
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt symmetric key: %w", err)
	}

	decrypted, err := sk.Decrypt([]byte(attachment.DataEnc))
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt attachment: %w", err)
	}

	return decrypted.GetBinary(), nil
}