Merge pull request #247 from Liam-Rougoor/feature/238-specify-public-…

…key-in-ring Rework to support usage of encryption subkeys
mattosaurus · Oct 18, 2023 · 74f1bae · 74f1bae
2 parents d6b0537 + 8e52de2
commit 74f1bae
Show file tree

Hide file tree

Showing 6 changed files with 157 additions and 9 deletions.
diff --git a/PgpCore.Tests/UnitTests/UnitTestsSync.cs b/PgpCore.Tests/UnitTests/UnitTestsSync.cs
@@ -3,8 +3,6 @@
 using System.Diagnostics;
 using System.IO;
 using System.Linq;
-using System.Security.Cryptography;
-using System.Threading.Tasks;
 using Org.BouncyCastle.Bcpg;
 using Org.BouncyCastle.Bcpg.OpenPgp;
 using Xunit;
@@ -358,7 +356,6 @@ public void Decrypt300MbFile_DecryptEncryptedFileWithMemoryUsageLessThan50Mb(Key
             // Arrange
             long memoryCap = 50 * 1024 * 1024;
             long startPeakWorkingSet = Process.GetCurrentProcess().PeakWorkingSet64;
-
             TestFactory testFactory = new TestFactory();
             testFactory.Arrange(keyType, fileType);
             EncryptionKeys encryptionKeys = new EncryptionKeys(testFactory.PublicKeyFileInfo);
@@ -1727,6 +1724,42 @@ public void DecryptStream_DecryptEncryptedStream(KeyType keyType)
             testFactory.Teardown();
         }
 
+        [Fact]
+        public void DecryptStream_DecryptEncryptedStreamWithPreferredKey()
+        {
+            // Arrange
+            TestFactory testFactory = new TestFactory();
+            testFactory.Arrange(KeyType.Known, FileType.Known);
+            EncryptionKeys encryptionKeys = new EncryptionKeys(testFactory.PublicKeyStream);
+            EncryptionKeys decryptionKeys = new EncryptionKeys(testFactory.PrivateKeyStream, testFactory.Password);
+            PGP pgpEncrypt = new PGP(encryptionKeys);
+            PGP pgpDecrypt = new PGP(decryptionKeys);
+
+            long[] keyIdsInPublicKeyRing = encryptionKeys.PublicKeyRings.First().PgpPublicKeyRing.GetPublicKeys()
+                .Where(key => key.IsEncryptionKey).Select(key => key.KeyId).ToArray();
+            foreach (long keyId in keyIdsInPublicKeyRing)
+            {
+                // Act
+                encryptionKeys.UseEncrytionKey(keyId);
+                using (Stream inputFileStream = testFactory.ContentStream)
+                using (Stream outputFileStream = File.Create(testFactory.EncryptedContentFilePath))
+                    pgpEncrypt.EncryptStream(inputFileStream, outputFileStream);
+
+                using (Stream inputFileStream = testFactory.EncryptedContentStream)
+                using (Stream outputFileStream = File.Create(testFactory.DecryptedContentFilePath))
+                    pgpDecrypt.DecryptStream(inputFileStream, outputFileStream);
+
+                // Assert
+                Assert.True(testFactory.EncryptedContentFileInfo.Exists);
+                Assert.True(testFactory.DecryptedContentFileInfo.Exists);
+                Assert.Equal(testFactory.Content, testFactory.DecryptedContent.Trim());
+            }
+            Assert.True(keyIdsInPublicKeyRing.Length > 1);
+
+            // Teardown
+            testFactory.Teardown();
+        }
+
         [Theory]
         [InlineData(KeyType.Generated)]
         [InlineData(KeyType.Known)]
@@ -1891,6 +1924,39 @@ public void Verify_VerifyEncryptedAndSignedStream(KeyType keyType)
             // Teardown
             testFactory.Teardown();
         }
+
+        [Fact]
+        public void Verify_VerifyEncryptedAndSignedStreamForMultipleKeys()
+        {
+            // Arrange
+            TestFactory testFactory = new TestFactory();
+            testFactory.Arrange(KeyType.Known, FileType.Known);
+            EncryptionKeys encryptionKeys = new EncryptionKeys(testFactory.PublicKeyStream, testFactory.PrivateKeyStream, testFactory.Password);
+            PGP pgp = new PGP(encryptionKeys);
+
+            // Act
+            long[] keyIdsInPublicKeyRing = encryptionKeys.PublicKeyRings.First().PgpPublicKeyRing.GetPublicKeys()
+                .Where(key => key.IsEncryptionKey).Select(key => key.KeyId).ToArray();
+            foreach (long keyId in keyIdsInPublicKeyRing)
+            {
+                encryptionKeys.UseEncrytionKey(keyId);
+                using (Stream inputFileStream = testFactory.ContentStream)
+                using (Stream outputFileStream = File.Create(testFactory.EncryptedContentFilePath))
+                    pgp.EncryptStreamAndSign(inputFileStream, outputFileStream);
+
+                bool verified = false;
+
+                using (Stream inputFileStream = testFactory.EncryptedContentStream)
+                    verified = pgp.VerifyStream(inputFileStream);
+
+                // Assert
+                Assert.True(testFactory.EncryptedContentFileInfo.Exists);
+                Assert.True(verified);
+            }
+
+            // Teardown
+            testFactory.Teardown();
+        }
 
         [Theory]
         [InlineData(KeyType.Generated)]

diff --git a/PgpCore/EncryptionKeys.cs b/PgpCore/EncryptionKeys.cs
@@ -10,6 +10,7 @@ public class EncryptionKeys : IEncryptionKeys
 	{
 		#region Instance Members (Public)
 
+		public IEnumerable<PgpPublicKeyRingWithPreferredKey> PublicKeyRings => _publicKeyRingsWithPreferredKey.Value;
 		public IEnumerable<PgpPublicKey> EncryptKeys => _encryptKeys.Value;
 		public IEnumerable<PgpPublicKey> VerificationKeys => _verificationKeys.Value;
 		public PgpPrivateKey SigningPrivateKey => _signingPrivateKey.Value;
@@ -33,6 +34,7 @@ public class EncryptionKeys : IEncryptionKeys
 		private Lazy<PgpPrivateKey> _signingPrivateKey;
 		private Lazy<PgpSecretKey> _signingSecretKey;
 		private Lazy<PgpSecretKeyRingBundle> _secretKeys;
+		private Lazy<IEnumerable<PgpPublicKeyRingWithPreferredKey>> _publicKeyRingsWithPreferredKey;
 
 		#endregion Instance Members (Private)
 
@@ -362,6 +364,19 @@ public PgpPrivateKey FindSecretKey(long keyId)
 			return pgpSecKey.ExtractPrivateKey(_passPhrase.ToCharArray());
 		}
 
+		/// <summary>
+		/// This method will try to find the key with the given keyId in a key ring and set it as the preferred key.
+		/// If it cannot find the key, it will not change the preferred key.
+		/// </summary>
+		/// <param name="keyId">The keyId to find.</param>
+		public void UseEncrytionKey(long keyId)
+		{
+			foreach (PgpPublicKeyRingWithPreferredKey publicKeyRing in PublicKeyRings)
+			{
+				publicKeyRing.UsePreferredEncryptionKey(keyId);
+			}
+		}
+
 		#endregion Public Methods
 
 		#region Private Key
@@ -389,10 +404,12 @@ private void
 				_masterKey = new Lazy<PgpPublicKey>(() => null);
 				_encryptKeys = new Lazy<IEnumerable<PgpPublicKey>>(() => null);
 				_verificationKeys = new Lazy<IEnumerable<PgpPublicKey>>(() => null);
+				_publicKeyRingsWithPreferredKey = new Lazy<IEnumerable<PgpPublicKeyRingWithPreferredKey>>(() => null);
 			}
 			else
 			{
 				// Need to consume the stream into a list before it is closed (can happen because of lazy instantiation).
+				_publicKeyRingsWithPreferredKey = new Lazy<IEnumerable<PgpPublicKeyRingWithPreferredKey>>(() => publicKeyRings.Select(keyRing => new PgpPublicKeyRingWithPreferredKey(keyRing)).ToArray());
 				_masterKey = new Lazy<PgpPublicKey>(() =>
 					Utilities.FindMasterKey(publicKeyRings.First()));
 				_encryptKeys = new Lazy<IEnumerable<PgpPublicKey>>(() =>

diff --git a/PgpCore/IEncryptionKeys.cs b/PgpCore/IEncryptionKeys.cs
@@ -12,6 +12,7 @@ namespace PgpCore
     /// </summary>
     public interface IEncryptionKeys
     {
+        IEnumerable<PgpPublicKeyRingWithPreferredKey> PublicKeyRings { get; }
         IEnumerable<PgpPublicKey> EncryptKeys { get; }
         IEnumerable<PgpPublicKey> VerificationKeys { get; }
         PgpPrivateKey SigningPrivateKey { get; }

diff --git a/PgpCore/PGP.cs b/PgpCore/PGP.cs
@@ -552,8 +552,9 @@ public async Task EncryptStreamAsync(Stream inputStream, Stream outputStream, bo
 
 			PgpEncryptedDataGenerator pk =
 				new PgpEncryptedDataGenerator(SymmetricKeyAlgorithm, withIntegrityCheck, new SecureRandom());
-			foreach (PgpPublicKey publicKey in EncryptionKeys.EncryptKeys)
+			foreach (PgpPublicKeyRingWithPreferredKey publicKeyRing in EncryptionKeys.PublicKeyRings)
 			{
+				PgpPublicKey publicKey = publicKeyRing.PreferredEncryptionKey ?? publicKeyRing.DefaultEncryptionKey;
 				pk.AddMethod(publicKey);
 			}
 
@@ -670,8 +671,9 @@ public void EncryptStream(Stream inputStream, Stream outputStream, bool armor =
 			PgpEncryptedDataGenerator pk =
 				new PgpEncryptedDataGenerator(SymmetricKeyAlgorithm, withIntegrityCheck, new SecureRandom());
 
-			foreach (PgpPublicKey publicKey in EncryptionKeys.EncryptKeys)
+			foreach (PgpPublicKeyRingWithPreferredKey publicKeyRing in EncryptionKeys.PublicKeyRings)
 			{
+				PgpPublicKey publicKey = publicKeyRing.PreferredEncryptionKey ?? publicKeyRing.DefaultEncryptionKey;
 				pk.AddMethod(publicKey);
 			}
 
@@ -5688,8 +5690,8 @@ private Task<VerificationResult> VerifyAsync(Stream inputStream, bool throwIfEnc
 				PgpPublicKeyEncryptedData publicKeyEncryptedData = Utilities.ExtractPublicKey(dataList);
 				var keyIdToVerify = publicKeyEncryptedData.KeyId;
 				// If we encounter an encrypted packet, verify with the encryption keys used instead
-				verified = Utilities.FindPublicKey(keyIdToVerify, EncryptionKeys.EncryptKeys, out PgpPublicKey _);
-
+				// TODO does this even make sense? maybe throw exception instead, or try to decrypt first
+				verified = Utilities.FindPublicKeyInKeyRings(keyIdToVerify, EncryptionKeys.PublicKeyRings.Select(keyRing => keyRing.PgpPublicKeyRing), out PgpPublicKey _);
 			}
 			else if (pgpObject is PgpOnePassSignatureList onePassSignatureList)
 			{
@@ -5835,7 +5837,8 @@ private VerificationResult Verify(Stream inputStream, bool throwIfEncrypted = fa
 				// Verify against public key ID and that of any sub keys
 
 				// If we encounter an encrypted packet, verify the encryption key used instead
-				if (Utilities.FindPublicKey(keyIdToVerify, EncryptionKeys.EncryptKeys, out PgpPublicKey _))
+				// TODO does this even make sense? maybe throw exception instead, or try to decrypt first
+				if (Utilities.FindPublicKeyInKeyRings(keyIdToVerify, EncryptionKeys.PublicKeyRings.Select(keyRing => keyRing.PgpPublicKeyRing), out PgpPublicKey _))
 				{
 					verified = true;
 				}
@@ -6148,8 +6151,9 @@ private Stream ChainEncryptedOut(Stream outputStream, bool withIntegrityCheck)
 			var encryptedDataGenerator =
 				new PgpEncryptedDataGenerator(SymmetricKeyAlgorithm, withIntegrityCheck, new SecureRandom());
 
-			foreach (PgpPublicKey publicKey in EncryptionKeys.EncryptKeys)
+			foreach (PgpPublicKeyRingWithPreferredKey publicKeyRing in EncryptionKeys.PublicKeyRings)
 			{
+				PgpPublicKey publicKey = publicKeyRing.PreferredEncryptionKey ?? publicKeyRing.DefaultEncryptionKey;
 				encryptedDataGenerator.AddMethod(publicKey);
 			}
 

diff --git a/PgpCore/PgpPublicKeyRingWithPreferredKey.cs b/PgpCore/PgpPublicKeyRingWithPreferredKey.cs
@@ -0,0 +1,45 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Org.BouncyCastle.Bcpg.OpenPgp;
+
+namespace PgpCore
+{
+    /// <summary>
+    /// A wrapper class for <see cref="PgpPublicKeyRing"/> that also keeps track of a preferred <see cref="PgpPublicKey"/> to be used for encryption.
+    /// </summary>
+    public class PgpPublicKeyRingWithPreferredKey
+    {
+        public PgpPublicKeyRing PgpPublicKeyRing { get; set; }
+        public PgpPublicKey PreferredEncryptionKey { get; private set; } = null;
+        public PgpPublicKey DefaultEncryptionKey => _defaultEncryptionKey.Value;
+
+        private Lazy<PgpPublicKey> _defaultEncryptionKey;
+        private Lazy<IEnumerable<PgpPublicKey>> _encryptionKeys;
+
+        public PgpPublicKeyRingWithPreferredKey(PgpPublicKeyRing publicKeyRing)
+        {
+            PgpPublicKeyRing = publicKeyRing;
+            _defaultEncryptionKey = new Lazy<PgpPublicKey>(() => Utilities.FindBestEncryptionKey(PgpPublicKeyRing));
+            _encryptionKeys = new Lazy<IEnumerable<PgpPublicKey>>(() => PgpPublicKeyRing.GetPublicKeys().Where(key => key.IsEncryptionKey));
+        }
+
+        /// <summary>
+        /// Try to find the key with the given keyId and set it as the preferred encryption key.
+        /// If no key is found, the preferred key is not changed.
+        /// </summary>
+        /// <param name="keyId">The keyId to find.</param>
+        public void UsePreferredEncryptionKey(long? keyId)
+        {
+            PreferredEncryptionKey = _encryptionKeys.Value.FirstOrDefault(key => key.KeyId == keyId) ?? PreferredEncryptionKey;
+        }
+
+        /// <summary>
+        /// Clear the preferred encryption key.
+        /// </summary>
+        public void ClearPreferredEncryptionKey()
+        {
+            PreferredEncryptionKey = null;
+        }
+    }
+}
diff --git a/PgpCore/Utilities.cs b/PgpCore/Utilities.cs
@@ -640,6 +640,21 @@ public static bool FindPublicKey(long keyId, IEnumerable<PgpPublicKey> verificat
 			return foundKeys.Any();
 		}
 
+		public static bool FindPublicKeyInKeyRings(long keyId, IEnumerable<PgpPublicKeyRing> publicKeyRings,
+			out PgpPublicKey verificationKey)
+		{
+			verificationKey = null;
+
+			foreach (PgpPublicKeyRing publicKeyRing in publicKeyRings)
+			{
+				var verificationKeys = publicKeyRing.GetPublicKeys();
+				if (FindPublicKey(keyId, verificationKeys, out verificationKey))
+					return true;
+			}
+
+			return false;
+		}
+
 		private static async Task PipeFileContentsAsync(FileInfo file, Stream pOut, int bufSize)
 		{
 			using (FileStream inputStream = file.OpenRead())