From 8d42a537753204abc81d3c2e99c9c08abd62f926 Mon Sep 17 00:00:00 2001
From: Ilya Chesnokov <chesnokov.ilya@protonmail.com>
Date: Thu, 2 Jul 2020 15:55:11 +0700
Subject: [PATCH] Add getEntity and getEcryptionKeyIDs functions to key and
 message types respectively (#55)

* add getEntity function to key struct

* add getEncryptionKeyIDs

* add chengelog + bool return in getEncryptionKeyIDs

* fix description
---
 CHANGELOG.md           | 10 ++++++++++
 crypto/key.go          |  5 +++++
 crypto/key_test.go     | 11 +++++++++++
 crypto/message.go      | 21 +++++++++++++++++++++
 crypto/message_test.go | 17 +++++++++++++++++
 5 files changed, 64 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 20d0da8..2b41e93 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 (msg *PGPMessage) GetArmoredWithCustomHeaders(comment, version string) (string, error)
 ```
 
+- Now one could get key IDs of keys to which the session key is encrypted
+```go
+func (msg *PGPMessage) getEncryptionKeyIDs() ([]uint64, bool)
+```
+
+get x/crypto Entity (components of an OpenPGP key) from Key struct
+```go
+(key *Key) GetEntity() *openpgp.Entity
+```
+
 ### Changed
 - Improved key and message armoring testing
 
diff --git a/crypto/key.go b/crypto/key.go
index d516316..bfbb0e1 100644
--- a/crypto/key.go
+++ b/crypto/key.go
@@ -354,6 +354,11 @@ func (key *Key) GetSHA256Fingerprints() (fingerprints []string) {
 	return
 }
 
+// GetEntity gets x/crypto Entity object.
+func (key *Key) GetEntity() *openpgp.Entity {
+	return key.entity
+}
+
 // --- Internal methods
 
 // getSHA256FingerprintBytes computes the SHA256 fingerprint of a public key
diff --git a/crypto/key_test.go b/crypto/key_test.go
index 3505587..25b1c86 100644
--- a/crypto/key_test.go
+++ b/crypto/key_test.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 	"testing"
 
+	"golang.org/x/crypto/openpgp"
 	"golang.org/x/crypto/openpgp/armor"
 	"golang.org/x/crypto/rsa"
 
@@ -378,3 +379,13 @@ func TestGetSHA256FingerprintsV4(t *testing.T) {
 	assert.Exactly(t, "d9ac0b857da6d2c8be985b251a9e3db31e7a1d2d832d1f07ebe838a9edce9c24", sha256Fingerprints[0])
 	assert.Exactly(t, "203dfba1f8442c17e59214d9cd11985bfc5cc8721bb4a71740dd5507e58a1a0d", sha256Fingerprints[1])
 }
+
+func TestGetEntity(t *testing.T) {
+	publicKey, err := NewKeyFromArmored(readTestFile("keyring_publicKey", false))
+	if err != nil {
+		t.Fatal("Cannot unarmor key:", err)
+	}
+	entity := publicKey.GetEntity()
+	assert.True(t, entity.PrimaryIdentity().SelfSignature.FlagsValid)
+	assert.IsType(t, &openpgp.Entity{}, entity)
+}
diff --git a/crypto/message.go b/crypto/message.go
index 1471fda..d3921b9 100644
--- a/crypto/message.go
+++ b/crypto/message.go
@@ -223,6 +223,27 @@ func (msg *PGPMessage) GetArmoredWithCustomHeaders(comment, version string) (str
 	return armor.ArmorWithTypeAndCustomHeaders(msg.Data, constants.PGPMessageHeader, version, comment)
 }
 
+// getEncryptionKeyIds Returns the key IDs of the keys to which the session key is encrypted.
+func (msg *PGPMessage) getEncryptionKeyIDs() ([]uint64, bool) {
+	packets := packet.NewReader(bytes.NewReader(msg.Data))
+	var err error
+	var ids []uint64
+	for {
+		var p packet.Packet
+		if p, err = packets.Next(); err == io.EOF {
+			break
+		}
+		enc, ok := p.(*packet.EncryptedKey)
+		if ok {
+			ids = append(ids, enc.KeyId)
+		}
+	}
+	if len(ids) > 0 {
+		return ids, true
+	}
+	return ids, false
+}
+
 // GetBinaryDataPacket returns the unarmored binary datapacket as a []byte.
 func (msg *PGPSplitMessage) GetBinaryDataPacket() []byte {
 	return msg.DataPacket
diff --git a/crypto/message_test.go b/crypto/message_test.go
index 9548d8d..e3cf53a 100644
--- a/crypto/message_test.go
+++ b/crypto/message_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/base64"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/crypto/openpgp/packet"
@@ -208,6 +209,22 @@ func TestMultipleKeyMessageEncryption(t *testing.T) {
 	assert.Exactly(t, message.GetString(), decrypted.GetString())
 }
 
+func TestMessagegetGetEncryptionKeyIDs(t *testing.T) {
+	var message = NewPlainMessageFromString("plain text")
+	assert.Exactly(t, 3, len(keyRingTestMultiple.entities))
+
+	ciphertext, err := keyRingTestMultiple.Encrypt(message, keyRingTestPrivate)
+	if err != nil {
+		t.Fatal("Expected no error when encrypting, got:", err)
+	}
+	ids, ok := ciphertext.getEncryptionKeyIDs()
+	assert.Exactly(t, 3, len(ids))
+	assert.True(t, ok)
+	encKey, ok := keyRingTestMultiple.entities[0].EncryptionKey(time.Now())
+	assert.True(t, ok)
+	assert.Exactly(t, encKey.PublicKey.KeyId, ids[0])
+}
+
 func TestMessageGetArmoredWithCustomHeaders(t *testing.T) {
 	var message = NewPlainMessageFromString("plain text")