package pgp

import (
	"bytes"
	"testing"

	"github.com/ProtonMail/go-crypto/openpgp"
	"github.com/ProtonMail/go-crypto/openpgp/armor"
)

// serializeUnencryptedPrivate serializes an as-yet-unprotected private key — only
// safe to call before EncryptPrivateKeys (see the gotcha documented on
// serializePrivateWithoutSigning in identity.go). No production code path needs
// this (every stored key goes through EncryptPrivateKeys first); it exists here
// purely to simulate a genuinely unprotected "gpg --export-secret-keys" output for
// TestImportPrivateKeyUnencrypted.
func serializeUnencryptedPrivate(entity *openpgp.Entity) ([]byte, error) {
	var buf bytes.Buffer
	w, err := armor.Encode(&buf, openpgp.PrivateKeyType, nil)
	if err != nil {
		return nil, err
	}
	if err := entity.SerializePrivate(w, nil); err != nil {
		return nil, err
	}
	if err := w.Close(); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func testEntity() Entity {
	return Entity{
		Headers: []string{"Content-Type: text/plain; charset=utf-8"},
		Body:    []byte("hello world\r\nsecond line\r\n"),
	}
}

func TestGenerateEncryptDecryptRoundTrip(t *testing.T) {
	pubPEM, privPEM, err := GenerateKeyPair("alice@example.com", "correct horse battery staple")
	if err != nil {
		t.Fatalf("GenerateKeyPair: %v", err)
	}
	if bytes.Contains(privPEM, []byte("correct horse")) {
		t.Fatal("stored private key armor should not contain the plaintext passphrase")
	}

	recipient, err := ParsePublicKey(pubPEM)
	if err != nil {
		t.Fatalf("ParsePublicKey: %v", err)
	}

	orig := testEntity()
	encrypted, err := EncryptEntity(orig, []*openpgp.Entity{recipient})
	if err != nil {
		t.Fatalf("EncryptEntity: %v", err)
	}
	if ct := headerValue(encrypted.Headers, "Content-Type"); ct == "" {
		t.Fatal("expected a Content-Type header on the encrypted entity")
	}
	if bytes.Contains(encrypted.Body, orig.Body) {
		t.Fatal("encrypted body should not contain the plaintext")
	}

	identity, err := ParsePrivateKey(privPEM)
	if err != nil {
		t.Fatalf("ParsePrivateKey: %v", err)
	}
	if err := UnlockPrivateKey(identity, "correct horse battery staple"); err != nil {
		t.Fatalf("UnlockPrivateKey: %v", err)
	}

	decrypted, err := DecryptEntity(encrypted, identity)
	if err != nil {
		t.Fatalf("DecryptEntity: %v", err)
	}
	if !bytes.Equal(decrypted.Body, orig.Body) {
		t.Fatalf("body mismatch: got %q want %q", decrypted.Body, orig.Body)
	}
	if headerValue(decrypted.Headers, "Content-Type") != headerValue(orig.Headers, "Content-Type") {
		t.Fatalf("header mismatch: got %v want %v", decrypted.Headers, orig.Headers)
	}
}

func TestUnlockPrivateKeyWrongPassphraseFails(t *testing.T) {
	_, privPEM, err := GenerateKeyPair("alice@example.com", "right passphrase")
	if err != nil {
		t.Fatal(err)
	}
	identity, err := ParsePrivateKey(privPEM)
	if err != nil {
		t.Fatal(err)
	}
	if err := UnlockPrivateKey(identity, "wrong passphrase"); err == nil {
		t.Fatal("expected the wrong passphrase to fail")
	}
}

func TestEncryptMultipleRecipientsBothCanDecrypt(t *testing.T) {
	senderPub, senderPriv, err := GenerateKeyPair("sender@example.com", "sender pass")
	if err != nil {
		t.Fatal(err)
	}
	recipPub, recipPriv, err := GenerateKeyPair("recipient@example.com", "recipient pass")
	if err != nil {
		t.Fatal(err)
	}
	senderPubEntity, _ := ParsePublicKey(senderPub)
	recipPubEntity, _ := ParsePublicKey(recipPub)

	orig := testEntity()
	encrypted, err := EncryptEntity(orig, []*openpgp.Entity{senderPubEntity, recipPubEntity})
	if err != nil {
		t.Fatal(err)
	}

	senderIdentity, _ := ParsePrivateKey(senderPriv)
	if err := UnlockPrivateKey(senderIdentity, "sender pass"); err != nil {
		t.Fatal(err)
	}
	senderCopy, err := DecryptEntity(encrypted, senderIdentity)
	if err != nil {
		t.Fatalf("sender DecryptEntity: %v", err)
	}
	if !bytes.Equal(senderCopy.Body, orig.Body) {
		t.Fatal("sender's own copy did not decrypt to the original body")
	}

	recipIdentity, _ := ParsePrivateKey(recipPriv)
	if err := UnlockPrivateKey(recipIdentity, "recipient pass"); err != nil {
		t.Fatal(err)
	}
	recipCopy, err := DecryptEntity(encrypted, recipIdentity)
	if err != nil {
		t.Fatalf("recipient DecryptEntity: %v", err)
	}
	if !bytes.Equal(recipCopy.Body, orig.Body) {
		t.Fatal("recipient's copy did not decrypt to the original body")
	}
}

func TestImportPrivateKeyUnencrypted(t *testing.T) {
	// Simulate a raw, not-yet-passphrase-protected export by generating a key and
	// serializing it before EncryptPrivateKeys is ever called.
	entity, err := openpgp.NewEntity("bob@example.com", "", "bob@example.com", defaultConfig())
	if err != nil {
		t.Fatal(err)
	}
	unprotectedArmor, err := serializeUnencryptedPrivate(entity)
	if err != nil {
		t.Fatal(err)
	}

	pubPEM, privPEM, err := ImportPrivateKey(unprotectedArmor, "new passphrase")
	if err != nil {
		t.Fatalf("ImportPrivateKey: %v", err)
	}

	identity, err := ParsePrivateKey(privPEM)
	if err != nil {
		t.Fatal(err)
	}
	if !identity.PrivateKey.Encrypted {
		t.Fatal("expected the imported key to be encrypted after import")
	}
	if err := UnlockPrivateKey(identity, "new passphrase"); err != nil {
		t.Fatalf("expected the new passphrase to unlock the imported key: %v", err)
	}

	recipient, err := ParsePublicKey(pubPEM)
	if err != nil {
		t.Fatal(err)
	}
	orig := testEntity()
	encrypted, err := EncryptEntity(orig, []*openpgp.Entity{recipient})
	if err != nil {
		t.Fatal(err)
	}
	decrypted, err := DecryptEntity(encrypted, identity)
	if err != nil {
		t.Fatalf("DecryptEntity after import: %v", err)
	}
	if !bytes.Equal(decrypted.Body, orig.Body) {
		t.Fatal("round trip through an imported unencrypted key failed")
	}
}

func TestImportPrivateKeyAlreadyEncrypted(t *testing.T) {
	_, existingArmor, err := GenerateKeyPair("carol@example.com", "original passphrase")
	if err != nil {
		t.Fatal(err)
	}

	if _, _, err := ImportPrivateKey(existingArmor, "wrong passphrase"); err == nil {
		t.Fatal("expected import with the wrong passphrase for an already-encrypted key to fail")
	}

	pubPEM, privPEM, err := ImportPrivateKey(existingArmor, "original passphrase")
	if err != nil {
		t.Fatalf("ImportPrivateKey with the correct passphrase: %v", err)
	}
	identity, err := ParsePrivateKey(privPEM)
	if err != nil {
		t.Fatal(err)
	}
	if err := UnlockPrivateKey(identity, "original passphrase"); err != nil {
		t.Fatalf("expected the original passphrase to still unlock after re-import: %v", err)
	}

	recipient, err := ParsePublicKey(pubPEM)
	if err != nil {
		t.Fatal(err)
	}
	orig := testEntity()
	encrypted, err := EncryptEntity(orig, []*openpgp.Entity{recipient})
	if err != nil {
		t.Fatal(err)
	}
	decrypted, err := DecryptEntity(encrypted, identity)
	if err != nil {
		t.Fatalf("DecryptEntity after re-import: %v", err)
	}
	if !bytes.Equal(decrypted.Body, orig.Body) {
		t.Fatal("round trip through a re-imported already-encrypted key failed")
	}
}

func TestFingerprintIsStableAndNonEmpty(t *testing.T) {
	pubPEM, _, err := GenerateKeyPair("alice@example.com", "pass")
	if err != nil {
		t.Fatal(err)
	}
	entity, err := ParsePublicKey(pubPEM)
	if err != nil {
		t.Fatal(err)
	}
	fp := Fingerprint(entity)
	if len(fp) == 0 {
		t.Fatal("expected a non-empty fingerprint")
	}
	entity2, err := ParsePublicKey(pubPEM)
	if err != nil {
		t.Fatal(err)
	}
	if Fingerprint(entity2) != fp {
		t.Fatal("expected the fingerprint to be stable across re-parses of the same key")
	}
}