Allowing cert to be exported even after loading.

Author: Nirjan Chapagain

src/VirtualClient/VirtualClient.Core/Identity/CertificateLoaderHelper.cs

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+namespace VirtualClient.Identity
+{
+    using System.Security.Cryptography.X509Certificates;
+
+    /// <summary>
+    /// Certificate Loader  to cleanly handle differences in .NET versions for loading certificates from byte arrays.
+    /// </summary>
+    internal static class CertificateLoaderHelper
+    {
+        internal static X509Certificate2 LoadPublic(byte[] cerBytes)
+        {
+#if NET9_0_OR_GREATER
+            return X509CertificateLoader.LoadCertificate(cerBytes);
+#else
+        return new X509Certificate2(cerBytes);
+#endif
+        }
+
+        internal static X509Certificate2 LoadPkcs12(
+            byte[] pfxBytes,
+            string password,
+            X509KeyStorageFlags flags)
+        {
+#if NET9_0_OR_GREATER
+            return X509CertificateLoader.LoadPkcs12(
+                pfxBytes,
+                password,
+                flags);
+#else
+        return new X509Certificate2(
+            pfxBytes,
+            password,
+            flags);
+#endif
+        }
+    }
+}

src/VirtualClient/VirtualClient.Core/KeyVaultManager.cs

@@ -15,6 +15,7 @@ namespace VirtualClient
     using Azure.Security.KeyVault.Secrets;
     using Polly;
     using VirtualClient.Common.Extensions;
+    using VirtualClient.Identity;
 
     /// <summary>
     /// Provides methods for retrieving secrets, keys, and certificates from an Azure Key Vault.
@@ -211,40 +212,50 @@ public async Task<X509Certificate2> GetCertificateAsync(
             this.StoreDescription.ThrowIfNull(nameof(this.StoreDescription));
             certName.ThrowIfNullOrWhiteSpace(nameof(certName), "The certificate name cannot be null or empty.");
 
-            // Use the keyVaultUri if provided as a parameter, otherwise use the store's EndpointUri
             Uri vaultUri = !string.IsNullOrWhiteSpace(keyVaultUri)
                 ? new Uri(keyVaultUri)
                 : ((DependencyKeyVaultStore)this.StoreDescription).EndpointUri;
 
             CertificateClient client = this.CreateCertificateClient(vaultUri, ((DependencyKeyVaultStore)this.StoreDescription).Credentials);
 
+            var credentials = ((DependencyKeyVaultStore)this.StoreDescription).Credentials;
+
+            CertificateClient certificateClient = this.CreateCertificateClient(vaultUri, credentials);
+            SecretClient secretClient = this.CreateSecretClient(vaultUri, credentials);
+
             try
             {
                 return await (retryPolicy ?? KeyVaultManager.DefaultRetryPolicy).ExecuteAsync(async () =>
                 {
                     // Get the full certificate with private key (PFX) if requested
                     if (retrieveWithPrivateKey)
                     {
-                        X509Certificate2 privateKeyCert = await client
-                            .DownloadCertificateAsync(certName, cancellationToken: cancellationToken)
-                            .ConfigureAwait(false);
+                        KeyVaultSecret secret = await secretClient.GetSecretAsync(certName, cancellationToken: cancellationToken);
+
+                        if (secret?.Value == null)
+                        {
+                            throw new DependencyException($"Secret for certificate '{certName}' not found in vault '{vaultUri}'.");
+                        }
+
+                        byte[] pfxBytes = Convert.FromBase64String(secret.Value);
 
-                        if (privateKeyCert is null || !privateKeyCert.HasPrivateKey)
+                        X509Certificate2 pfxCertificate = CertificateLoaderHelper.LoadPkcs12(
+                            pfxBytes,
+                            string.Empty,
+                            X509KeyStorageFlags.Exportable | X509KeyStorageFlags.PersistKeySet);
+
+                        if (!pfxCertificate.HasPrivateKey)
                         {
-                            throw new DependencyException("Failed to retrieve certificate content with private key.");
+                            throw new DependencyException($"Certificate '{certName}' does not contain a private key.");
                         }
 
-                        return privateKeyCert;
+                        return pfxCertificate;
                     }
                     else
                     {
-                        // If private key not needed, load cert from PublicBytes
-                        KeyVaultCertificateWithPolicy cert = await client.GetCertificateAsync(certName, cancellationToken: cancellationToken);
-#if NET9_0_OR_GREATER
-                        return X509CertificateLoader.LoadCertificate(cert.Cer);
-#elif NET8_0_OR_GREATER
-                        return new X509Certificate2(cert.Cer);
-#endif
+                        // Public certificate only
+                        KeyVaultCertificateWithPolicy certBundle = await certificateClient.GetCertificateAsync(certName, cancellationToken: cancellationToken);
+                        return CertificateLoaderHelper.LoadPublic(certBundle.Cer);
                     }
                 }).ConfigureAwait(false);
             }
@@ -269,13 +280,6 @@ public async Task<X509Certificate2> GetCertificateAsync(
                     ex,
                     ErrorReason.HttpNonSuccessResponse);
             }
-            catch (Exception ex)
-            {
-                throw new DependencyException(
-                    $"Failed to get certificate '{certName}' from vault '{vaultUri}'.",
-                    ex,
-                    ErrorReason.HttpNonSuccessResponse);
-            }
         }
 
         /// <summary>
@@ -328,4 +332,4 @@ private void ValidateKeyVaultStore()
             }
         }
     }
-}
+}

src/VirtualClient/VirtualClient.Dependencies/CertificateInstallation.cs

@@ -62,7 +62,7 @@ public string CertificateName
         /// <summary>
         /// Gets the path to the file where the access token is saved.
         /// </summary>
-        public string AccessTokenPath 
+        public string AccessTokenPath
         {
             get
             {
@@ -138,12 +138,16 @@ protected override async Task ExecuteAsync(EventContext telemetryContext, Cancel
                     throw new PlatformNotSupportedException($"The '{nameof(CertificateInstallation)}' component is not supported on platform '{this.Platform}'.");
                 }
 
-                // If a download directory is specified, we will also export the certificate to that location.
+                // Export the certificate if requested
                 if (!string.IsNullOrEmpty(this.CertificateDownloadDir))
                 {
-                    string certificateFileName = this.WithPrivateKey 
+                    string certificateFileName = this.WithPrivateKey
                         ? $"{this.CertificateName}.pfx"
                         : $"{this.CertificateName}.cer";
+                    
+                    X509ContentType contentType = this.WithPrivateKey 
+                        ? X509ContentType.Pfx 
+                        : X509ContentType.Cert;
 
                     string certificatePath = this.Combine(this.CertificateDownloadDir, certificateFileName);
 
@@ -153,9 +157,8 @@ protected override async Task ExecuteAsync(EventContext telemetryContext, Cancel
                         this.fileSystem.File.Delete(certificatePath);
                     }
 
-                    // Export the new certificate
-                    await this.fileSystem.File.WriteAllBytesAsync(certificatePath, certificate.Export(X509ContentType.Pfx));
-                    Console.WriteLine($"Certificate exported to {certificatePath}");
+                    byte[] certBytes = certificate.Export(contentType, string.Empty);
+                    await this.fileSystem.File.WriteAllBytesAsync(certificatePath, certBytes);
                 }
             }
             catch (Exception exc)
@@ -276,5 +279,22 @@ protected IKeyVaultManager GetKeyVaultManager()
                     $"Either valid --KeyVault or --Token or --TokenPath must be passed in order to set up authentication with Key Vault.");
             }
         }
+
+        /// <summary>
+        /// Tries to get certificate data, returning null if an exception occurs.
+        /// </summary>
+        private byte[] TryGetCertData(Func<byte[]> getCertData)
+        {
+            try
+            {
+                return getCertData();
+            }
+            catch (Exception exc)
+            {
+                Console.WriteLine(exc.ToString());
+                Console.WriteLine("\n\n\n\n=================================================================================\n");
+                return null;
+            }
+        }
     }
 }