From e4da63c9f7124207fd59b25590e75a83969fe570 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sun, 29 Oct 2023 22:24:04 +0800
Subject: [PATCH] feat: updates

---
 .../hatter/tool/signpdf/CreateSignature.java  |  3 +-
 .../tool/signpdf/CreateSignatureBase.java     | 38 +++++++++++--------
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/src/main/java/me/hatter/tool/signpdf/CreateSignature.java b/src/main/java/me/hatter/tool/signpdf/CreateSignature.java
index 7ec22d0..72be307 100644
--- a/src/main/java/me/hatter/tool/signpdf/CreateSignature.java
+++ b/src/main/java/me/hatter/tool/signpdf/CreateSignature.java
@@ -13,6 +13,7 @@ import java.security.PrivateKey;
 import java.security.UnrecoverableKeyException;
 import java.security.cert.Certificate;
 import java.security.cert.CertificateException;
+import java.security.cert.X509Certificate;
 import java.util.Calendar;
 
 // https://svn.apache.org/viewvc/pdfbox/trunk/examples/src/main/
@@ -32,7 +33,7 @@ import java.util.Calendar;
 public class CreateSignature extends CreateSignatureBase {
     private final SignOptions signOptions;
 
-    public CreateSignature(SignOptions signOptions, Certificate[] certificateChain, PrivateKey privateKey)
+    public CreateSignature(SignOptions signOptions, X509Certificate[] certificateChain, PrivateKey privateKey)
             throws KeyStoreException, UnrecoverableKeyException, NoSuchAlgorithmException, CertificateException, IOException {
         super(certificateChain, privateKey);
         this.signOptions = signOptions;
diff --git a/src/main/java/me/hatter/tool/signpdf/CreateSignatureBase.java b/src/main/java/me/hatter/tool/signpdf/CreateSignatureBase.java
index bbd2688..fa45992 100644
--- a/src/main/java/me/hatter/tool/signpdf/CreateSignatureBase.java
+++ b/src/main/java/me/hatter/tool/signpdf/CreateSignatureBase.java
@@ -14,20 +14,21 @@ import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;
 import java.io.IOException;
 import java.io.InputStream;
 import java.security.*;
-import java.security.cert.Certificate;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
+import java.security.interfaces.ECPublicKey;
+import java.security.interfaces.RSAPublicKey;
 import java.util.Arrays;
 
 // https://svn.apache.org/viewvc/pdfbox/trunk/examples/src/main/java/
 // org/apache/pdfbox/examples/signature/CreateSignatureBase.java?view=co
 public abstract class CreateSignatureBase implements SignatureInterface {
     private PrivateKey privateKey;
-    private Certificate[] certificateChain;
+    private X509Certificate[] certificateChain;
     private String tsaUrl;
     private boolean externalSigning;
 
-    public CreateSignatureBase(Certificate[] certificateChain, PrivateKey privateKey)
+    public CreateSignatureBase(X509Certificate[] certificateChain, PrivateKey privateKey)
             throws KeyStoreException, UnrecoverableKeyException, NoSuchAlgorithmException, IOException, CertificateException {
         // grabs the first alias from the keystore and get the private key. An
         // alternative method or constructor could be used for setting a specific
@@ -35,23 +36,20 @@ public abstract class CreateSignatureBase implements SignatureInterface {
         this.privateKey = privateKey;
         this.certificateChain = certificateChain;
 
-        final Certificate cert = certificateChain[0];
-        if (cert instanceof X509Certificate) {
-            // avoid expired certificate
-            ((X509Certificate) cert).checkValidity();
-            SigUtils.checkCertificateUsage((X509Certificate) cert);
-        }
+        final X509Certificate cert = certificateChain[0];
+        cert.checkValidity();
+        SigUtils.checkCertificateUsage(cert);
     }
 
     public final void setPrivateKey(PrivateKey privateKey) {
         this.privateKey = privateKey;
     }
 
-    public final void setCertificateChain(final Certificate[] certificateChain) {
+    public final void setCertificateChain(final X509Certificate[] certificateChain) {
         this.certificateChain = certificateChain;
     }
 
-    public Certificate[] getCertificateChain() {
+    public X509Certificate[] getCertificateChain() {
         return certificateChain;
     }
 
@@ -61,17 +59,25 @@ public abstract class CreateSignatureBase implements SignatureInterface {
 
     @Override
     public byte[] sign(InputStream content) throws IOException {
-        // cannot be done private (interface)
         try {
             final CMSSignedDataGenerator gen = new CMSSignedDataGenerator();
-            final X509Certificate cert = (X509Certificate) certificateChain[0];
+            final X509Certificate cert = certificateChain[0];
             // TODO use customer signer
-            final String signatureAlgorithm = "SHA256WithECDSA";
-//            final String  signatureAlgorithm = "SHA256WithRSA";
+            final boolean isEcPublicKey = cert.getPublicKey() instanceof ECPublicKey;
+            final boolean isRsaPublicKey = cert.getPublicKey() instanceof RSAPublicKey;
+
+            final String signatureAlgorithm;
+            if (isEcPublicKey) {
+                signatureAlgorithm = "SHA256WithECDSA";
+            } else if (isRsaPublicKey) {
+                signatureAlgorithm = "SHA256WithRSA";
+            } else {
+                throw new RuntimeException("Supported algorithm: " + cert.getPublicKey().getClass());
+            }
             final ContentSigner contentSigner = new JcaContentSignerBuilder(signatureAlgorithm).build(privateKey);
             gen.addSignerInfoGenerator(new JcaSignerInfoGeneratorBuilder(
                     new JcaDigestCalculatorProviderBuilder().build()).build(contentSigner, cert));
-            
+
             gen.addCertificates(new JcaCertStore(Arrays.asList(certificateChain)));
             final CMSProcessableInputStream msg = new CMSProcessableInputStream(content);
             CMSSignedData signedData = gen.generate(msg, false);