pop/internal/pgp/pgp.go

package pgp

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	openpgp "github.com/ProtonMail/gopenpgp/v2/crypto"
	"github.com/frenocorp/pop/internal/config"
)

// ExternalKey represents an imported external PGP key.
type ExternalKey struct {
	KeyID        string `json:"key_id"`
	Fingerprint  string `json:"fingerprint"`
	Emails       []string `json:"emails"`
	CreatedAt    string `json:"created_at"`
	ExpiresAt    string `json:"expires_at,omitempty"`
	IsSubkey     bool   `json:"is_subkey"`
	CanEncrypt   bool   `json:"can_encrypt"`
	CanSign      bool   `json:"can_sign"`
	TrustLevel   string `json:"trust_level"`
	ArmorFile    string `json:"armor_file"`
}

// KeyStore manages external PGP keys.
type KeyStore struct {
	configDir string
	keysDir   string
	keysFile  string
}

// NewKeyStore creates a new PGP key store.
func NewKeyStore() (*KeyStore, error) {
	cfg := config.NewConfigManager()
	configDir := cfg.ConfigDir()

	keysDir := filepath.Join(configDir, "pgp_keys")
	if err := os.MkdirAll(keysDir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create PGP keys directory: %w", err)
	}

	return &KeyStore{
		configDir: configDir,
		keysDir:   keysDir,
		keysFile:  filepath.Join(configDir, "pgp_keys.json"),
	}, nil
}

// ImportKey imports an external PGP key from armored ASCII text.
func (ks *KeyStore) ImportKey(armor string, trustLevel string) (*ExternalKey, error) {
	if trustLevel == "" {
		trustLevel = "unknown"
	}

	pgpKey, err := openpgp.NewKeyFromArmored(armor)
	if err != nil {
		return nil, fmt.Errorf("failed to parse PGP key: %w", err)
	}

	fingerprint := pgpKey.GetFingerprint()
	keyID := fingerprint[len(fingerprint)-8:]

	emails := []string{}
	if entity := pgpKey.GetEntity(); entity != nil {
		for _, uid := range entity.Identities {
			if uid != nil && uid.UserId != nil {
				emails = append(emails, uid.UserId.Email)
			}
		}
	}

	expiresAt := ""
	if pgpKey.GetEntity() != nil {
		expiresAt = time.Unix(int64(pgpKey.GetEntity().PrimaryKey.CreationTime.Unix()), 0).UTC().Format(time.RFC3339)
	}

	key := &ExternalKey{
		KeyID:       keyID,
		Fingerprint: fingerprint,
		Emails:      emails,
		CreatedAt:   time.Now().UTC().Format(time.RFC3339),
		ExpiresAt:   expiresAt,
		IsSubkey:    false,
		CanEncrypt:  pgpKey.CanEncrypt(),
		CanSign:     pgpKey.IsPrivate(),
		TrustLevel:  trustLevel,
		ArmorFile:   filepath.Join(ks.keysDir, keyID+".asc"),
	}

	if err := os.WriteFile(key.ArmorFile, []byte(armor), 0600); err != nil {
		return nil, fmt.Errorf("failed to write key file: %w", err)
	}

	if err := ks.saveKeyMetadata(key); err != nil {
		os.Remove(key.ArmorFile)
		return nil, err
	}

	return key, nil
}

// ImportKeyFromFile imports a PGP key from a file containing armored ASCII.
func (ks *KeyStore) ImportKeyFromFile(filePath string, trustLevel string) (*ExternalKey, error) {
	data, err := os.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read key file: %w", err)
	}
	return ks.ImportKey(string(data), trustLevel)
}

// ListKeys returns all imported external keys.
func (ks *KeyStore) ListKeys() ([]ExternalKey, error) {
	data, err := os.ReadFile(ks.keysFile)
	if err != nil {
		if os.IsNotExist(err) {
			return []ExternalKey{}, nil
		}
		return nil, fmt.Errorf("failed to read keys metadata: %w", err)
	}

	var keys []ExternalKey
	if err := json.Unmarshal(data, &keys); err != nil {
		return nil, fmt.Errorf("failed to parse keys metadata: %w", err)
	}

	return keys, nil
}

// GetKey retrieves a key by key ID or fingerprint.
func (ks *KeyStore) GetKey(identifier string) (*ExternalKey, error) {
	keys, err := ks.ListKeys()
	if err != nil {
		return nil, err
	}

	for _, key := range keys {
		if key.KeyID == identifier || key.Fingerprint == identifier {
			return &key, nil
		}
	}

	return nil, fmt.Errorf("key %q not found", identifier)
}

// RemoveKey removes an external key by key ID or fingerprint.
func (ks *KeyStore) RemoveKey(identifier string) error {
	keys, err := ks.ListKeys()
	if err != nil {
		return err
	}

	var keyToRemove *ExternalKey
	newKeys := make([]ExternalKey, 0, len(keys))
	for i := range keys {
		if keys[i].KeyID == identifier || keys[i].Fingerprint == identifier {
			keyToRemove = &keys[i]
			continue
		}
		newKeys = append(newKeys, keys[i])
	}

	if keyToRemove == nil {
		return fmt.Errorf("key %q not found", identifier)
	}

	if keyToRemove.ArmorFile != "" {
		os.Remove(keyToRemove.ArmorFile)
	}

	return ks.writeKeysMetadata(newKeys)
}

// GetKeyArmor returns the armored ASCII representation of a key.
func (ks *KeyStore) GetKeyArmor(identifier string) (string, error) {
	key, err := ks.GetKey(identifier)
	if err != nil {
		return "", err
	}

	data, err := os.ReadFile(key.ArmorFile)
	if err != nil {
		return "", fmt.Errorf("failed to read key armor: %w", err)
	}

	return string(data), nil
}

// EncryptData encrypts plaintext using a public key.
func (ks *KeyStore) EncryptData(identifier, plaintext string) (string, error) {
	armor, err := ks.GetKeyArmor(identifier)
	if err != nil {
		return "", err
	}

	key, err := openpgp.NewKeyFromArmored(armor)
	if err != nil {
		return "", fmt.Errorf("failed to parse key for encryption: %w", err)
	}

	if !key.CanEncrypt() {
		return "", fmt.Errorf("key %s cannot be used for encryption", identifier)
	}

	plainMsg := openpgp.NewPlainMessageFromString(plaintext)
	keyRing, err := openpgp.NewKeyRing(key)
	if err != nil {
		return "", fmt.Errorf("failed to create keyring: %w", err)
	}

	encryptedMsg, err := keyRing.Encrypt(plainMsg, keyRing)
	if err != nil {
		return "", fmt.Errorf("failed to encrypt data: %w", err)
	}

	result, err := encryptedMsg.GetArmored()
	if err != nil {
		return "", fmt.Errorf("failed to get armored encrypted data: %w", err)
	}

	return result, nil
}

// SignData signs plaintext using a private key.
func (ks *KeyStore) SignData(identifier, plaintext, passphrase string) (string, error) {
	armor, err := ks.GetKeyArmor(identifier)
	if err != nil {
		return "", err
	}

	key, err := openpgp.NewKeyFromArmored(armor)
	if err != nil {
		return "", fmt.Errorf("failed to parse key for signing: %w", err)
	}

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

	if !key.IsPrivate() {
		return "", fmt.Errorf("key %s cannot be used for signing", identifier)
	}

	keyRing, err := openpgp.NewKeyRing(key)
	if err != nil {
		return "", fmt.Errorf("failed to create keyring: %w", err)
	}

	plainMsg := openpgp.NewPlainMessageFromString(plaintext)
	signedMsg, err := keyRing.SignDetached(plainMsg)
	if err != nil {
		return "", fmt.Errorf("failed to sign data: %w", err)
	}

	result, err := signedMsg.GetArmored()
	if err != nil {
		return "", fmt.Errorf("failed to get armored signature: %w", err)
	}

	return result, nil
}

// DecryptData decrypts PGP-encrypted data using a private key.
func (ks *KeyStore) DecryptData(identifier, encryptedData, passphrase string) (string, error) {
	armor, err := ks.GetKeyArmor(identifier)
	if err != nil {
		return "", err
	}

	key, err := openpgp.NewKeyFromArmored(armor)
	if err != nil {
		return "", fmt.Errorf("failed to parse key for decryption: %w", err)
	}

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

	keyRing, err := openpgp.NewKeyRing(key)
	if err != nil {
		return "", fmt.Errorf("failed to create keyring: %w", err)
	}

	encryptedMsg, parseErr := openpgp.NewPGPMessageFromArmored(encryptedData)
	if parseErr != nil {
		return "", fmt.Errorf("failed to parse encrypted message: %w", parseErr)
	}

	plainMessage, err := keyRing.Decrypt(encryptedMsg, nil, 0)
	if err != nil {
		return "", fmt.Errorf("failed to decrypt data: %w", err)
	}

	return string(plainMessage.Data), nil
}

// VerifySignature verifies a detached signature.
func (ks *KeyStore) VerifySignature(keyID, message, signature string) (bool, error) {
	armor, err := ks.GetKeyArmor(keyID)
	if err != nil {
		return false, err
	}

	key, err := openpgp.NewKeyFromArmored(armor)
	if err != nil {
		return false, fmt.Errorf("failed to parse key for verification: %w", err)
	}

	sig, err := openpgp.NewPGPSignatureFromArmored(signature)
	if err != nil {
		return false, fmt.Errorf("failed to parse signature: %w", err)
	}

	keyRing, err := openpgp.NewKeyRing(key)
	if err != nil {
		return false, fmt.Errorf("failed to create keyring: %w", err)
	}

	plainMsg := openpgp.NewPlainMessage([]byte(message))
		err = keyRing.VerifyDetached(plainMsg, sig, 0)
		if err != nil {
			return false, fmt.Errorf("signature verification failed: %w", err)
		}

		return true, nil
}

// TrustKey sets the trust level for a key.
func (ks *KeyStore) TrustKey(identifier, trustLevel string) error {
	keys, err := ks.ListKeys()
	if err != nil {
		return err
	}

	found := false
	for i, key := range keys {
		if key.KeyID == identifier || key.Fingerprint == identifier {
			keys[i].TrustLevel = trustLevel
			found = true
			break
		}
	}

	if !found {
		return fmt.Errorf("key %q not found", identifier)
	}

	return ks.writeKeysMetadata(keys)
}

// ExportKey exports a key to a file.
func (ks *KeyStore) ExportKey(identifier, outputPath string) error {
	armor, err := ks.GetKeyArmor(identifier)
	if err != nil {
		return err
	}

	if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
		return fmt.Errorf("failed to create output directory: %w", err)
	}

	if err := os.WriteFile(outputPath, []byte(armor), 0600); err != nil {
		return fmt.Errorf("failed to write exported key: %w", err)
	}

	return nil
}

// GetKeyFingerprint returns the fingerprint of a key from armored data.
func GetKeyFingerprint(armor string) (string, error) {
	key, err := openpgp.NewKeyFromArmored(armor)
	if err != nil {
		return "", err
	}
	return key.GetFingerprint(), nil
}

// ParseKeyInfo extracts key information from armored PGP data without importing.
func ParseKeyInfo(armor string) (*ExternalKey, error) {
	key, err := openpgp.NewKeyFromArmored(armor)
	if err != nil {
		return nil, fmt.Errorf("failed to parse key: %w", err)
	}

	fingerprint := key.GetFingerprint()
	keyID := fingerprint[len(fingerprint)-8:]

	emails := []string{}
	if entity := key.GetEntity(); entity != nil {
		for _, uid := range entity.Identities {
			if uid != nil && uid.UserId != nil {
				emails = append(emails, uid.UserId.Email)
			}
		}
	}

	expiresAt := ""
	if entity := key.GetEntity(); entity != nil {
		expiresAt = time.Unix(int64(entity.PrimaryKey.CreationTime.Unix()), 0).UTC().Format(time.RFC3339)
	}

	return &ExternalKey{
		KeyID:       keyID,
		Fingerprint: fingerprint,
		Emails:      emails,
		CreatedAt:   expiresAt,
		ExpiresAt:   expiresAt,
		IsSubkey:    false,
		CanEncrypt:  key.CanEncrypt(),
		CanSign:     key.IsPrivate(),
	}, nil
}

func (ks *KeyStore) saveKeyMetadata(key *ExternalKey) error {
	keys, err := ks.ListKeys()
	if err != nil {
		return err
	}

	for _, k := range keys {
		if k.KeyID == key.KeyID {
			return fmt.Errorf("key with ID %s already imported", key.KeyID)
		}
	}

	keys = append(keys, *key)
	return ks.writeKeysMetadata(keys)
}

func (ks *KeyStore) writeKeysMetadata(keys []ExternalKey) error {
	data, err := json.MarshalIndent(keys, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal keys metadata: %w", err)
	}

	return os.WriteFile(ks.keysFile, data, 0600)
}

// KeyFromReader reads PGP key data from an io.Reader (useful for stdin).
func (ks *KeyStore) KeyFromReader(reader io.Reader, trustLevel string) (*ExternalKey, error) {
	data, err := io.ReadAll(reader)
	if err != nil {
		return nil, fmt.Errorf("failed to read key data: %w", err)
	}

	content := string(data)
	if !strings.Contains(content, "BEGIN PGP PUBLIC KEY BLOCK") &&
		!strings.Contains(content, "BEGIN PGP PRIVATE KEY BLOCK") {
		return nil, fmt.Errorf("input does not appear to be a PGP key (missing armor header)")
	}

	return ks.ImportKey(content, trustLevel)
}