From 987923f9d8fe00c15ca1ecb93fa2581dbe664f44 Mon Sep 17 00:00:00 2001
From: Lukas Burkhalter <lubux@users.noreply.github.com>
Date: Mon, 28 Aug 2023 10:31:40 +0200
Subject: [PATCH] Add quick check for session key decryption (#249)

This commit adds the function QuickCheckDecrypt to the helper package.
The function allows to check with high probability if a session key can decrypt a data packet given its 24-byte prefix.
It only works for SEIPDv1 data packets that uses AES as a cipher.
---
 helper/decrypt_check.go      | 86 ++++++++++++++++++++++++++++++++++++
 helper/decrypt_check_test.go | 43 ++++++++++++++++++
 2 files changed, 129 insertions(+)
 create mode 100644 helper/decrypt_check.go
 create mode 100644 helper/decrypt_check_test.go

diff --git a/helper/decrypt_check.go b/helper/decrypt_check.go
new file mode 100644
index 00000000..eeff1089
--- /dev/null
+++ b/helper/decrypt_check.go
@@ -0,0 +1,86 @@
+package helper
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"io"
+
+	"github.com/ProtonMail/go-crypto/openpgp/packet"
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
+	"github.com/pkg/errors"
+)
+
+const AES_BLOCK_SIZE = 16
+
+func supported(cipher packet.CipherFunction) bool {
+	switch cipher {
+	case packet.CipherAES128, packet.CipherAES192, packet.CipherAES256:
+		return true
+	case packet.CipherCAST5, packet.Cipher3DES:
+		return false
+	}
+	return false
+}
+
+func blockSize(cipher packet.CipherFunction) int {
+	switch cipher {
+	case packet.CipherAES128, packet.CipherAES192, packet.CipherAES256:
+		return AES_BLOCK_SIZE
+	case packet.CipherCAST5, packet.Cipher3DES:
+		return 0
+	}
+	return 0
+}
+
+func blockCipher(cipher packet.CipherFunction, key []byte) (cipher.Block, error) {
+	switch cipher {
+	case packet.CipherAES128, packet.CipherAES192, packet.CipherAES256:
+		return aes.NewCipher(key)
+	case packet.CipherCAST5, packet.Cipher3DES:
+		return nil, errors.New("gopenpgp: cipher not supported for quick check")
+	}
+	return nil, errors.New("gopenpgp: unknown cipher")
+}
+
+// QuickCheckDecryptReader checks with high probability if the provided session key
+// can decrypt a data packet given its 24 byte long prefix.
+// The method reads up to but not exactly 24 bytes from the prefixReader.
+// NOTE: Only works for SEIPDv1 packets with AES.
+func QuickCheckDecryptReader(sessionKey *crypto.SessionKey, prefixReader crypto.Reader) (bool, error) {
+	algo, err := sessionKey.GetCipherFunc()
+	if err != nil {
+		return false, errors.New("gopenpgp: cipher algorithm not found")
+	}
+	if !supported(algo) {
+		return false, errors.New("gopenpgp: cipher not supported for quick check")
+	}
+	packetParser := packet.NewReader(prefixReader)
+	_, err = packetParser.Next()
+	if err != nil {
+		return false, errors.New("gopenpgp: failed to parse packet prefix")
+	}
+
+	blockSize := blockSize(algo)
+	encryptedData := make([]byte, blockSize+2)
+	_, err = io.ReadFull(prefixReader, encryptedData)
+	if err != nil {
+		return false, errors.New("gopenpgp: prefix is too short to check")
+	}
+
+	blockCipher, err := blockCipher(algo, sessionKey.Key)
+	if err != nil {
+		return false, errors.New("gopenpgp: failed to initialize the cipher")
+	}
+	_ = packet.NewOCFBDecrypter(blockCipher, encryptedData, packet.OCFBNoResync)
+	return encryptedData[blockSize-2] == encryptedData[blockSize] &&
+		encryptedData[blockSize-1] == encryptedData[blockSize+1], nil
+}
+
+// QuickCheckDecrypt checks with high probability if the provided session key
+// can decrypt the encrypted data packet given its 24 byte long prefix.
+// The method only considers the first 24 bytes of the prefix slice (prefix[:24]).
+// NOTE: Only works for SEIPDv1 packets with AES.
+func QuickCheckDecrypt(sessionKey *crypto.SessionKey, prefix []byte) (bool, error) {
+	return QuickCheckDecryptReader(sessionKey, bytes.NewReader(prefix))
+}
diff --git a/helper/decrypt_check_test.go b/helper/decrypt_check_test.go
new file mode 100644
index 00000000..f7341114
--- /dev/null
+++ b/helper/decrypt_check_test.go
@@ -0,0 +1,43 @@
+package helper
+
+import (
+	"encoding/hex"
+	"testing"
+
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
+)
+
+const testQuickCheckSessionKey = `038c9cb9d408074e36bac22c6b90973082f86e5b01f38b787da3927000365a81`
+const testQuickCheckSessionKeyAlg = "aes256"
+const testQuickCheckDataPacket = `d2540152ab2518950f282d98d901eb93c00fb55a3bb30b3b517d6a356f57884bac6963060ebb167ffc3296e5e99ec058aeff5003a4784a0734a62861ae56d2921b9b790d50586cd21cad45e2d84ac93fb5d8af2ce6c5`
+
+func TestCheckDecrypt(t *testing.T) {
+	sessionKeyData, err := hex.DecodeString(testQuickCheckSessionKey)
+	if err != nil {
+		t.Error(err)
+	}
+	dataPacket, err := hex.DecodeString(testQuickCheckDataPacket)
+	if err != nil {
+		t.Error(err)
+	}
+	sessionKey := &crypto.SessionKey{
+		Key:  sessionKeyData,
+		Algo: testQuickCheckSessionKeyAlg,
+	}
+	ok, err := QuickCheckDecrypt(sessionKey, dataPacket[:22])
+	if err != nil {
+		t.Error(err)
+	}
+	if !ok {
+		t.Error("should be able to decrypt")
+	}
+
+	sessionKey.Key[0] += 1
+	ok, err = QuickCheckDecrypt(sessionKey, dataPacket[:22])
+	if err != nil {
+		t.Error(err)
+	}
+	if ok {
+		t.Error("should no be able to decrypt")
+	}
+}