
package de.ivu.elect.nl.x509;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.function.Consumer;

import sun.security.pkcs.ContentInfo;
import sun.security.pkcs.PKCS7;
import sun.security.pkcs.SignerInfo;
import sun.security.util.DerOutputStream;
import sun.security.util.DerValue;
import sun.security.x509.AlgorithmId;
import sun.security.x509.X500Name;

public class NlCMS {

  private static final String ALGORITHM = "Sha256WithRSA";

  public NlCMS() {
    super();
  }

  public void sign(OutputStream os, PrivateKey privateKey, PublicKey publicKey, NlX509Metadata nlX509Metadata, byte[] dataToSign) throws Exception {
    X509Certificate x509Certificate = getX509Certificate(privateKey, publicKey, nlX509Metadata);

    // Compute signature
    Signature signature = Signature.getInstance(ALGORITHM);
    signature.initSign(privateKey);
    signature.update(dataToSign);

    byte[] signatureData = signature.sign();

    // Load X500Name
    X500Name xName = X500Name.asX500Name(x509Certificate.getSubjectX500Principal());

    // Load serial number
    BigInteger serial = x509Certificate.getSerialNumber();

    // Load digest algorithm
    AlgorithmId digestAlgorithmId = new AlgorithmId(AlgorithmId.SHA256_oid);

    // Load signing algorithm
    AlgorithmId signAlgorithmId = new AlgorithmId(AlgorithmId.RSAEncryption_oid);

    // Create SignerInfo
    SignerInfo sInfo = new SignerInfo(xName, serial, digestAlgorithmId, signAlgorithmId, signatureData);

    // Create ContentInfo
    ContentInfo cInfo = new ContentInfo(ContentInfo.DATA_OID, new DerValue(DerValue.tag_OctetString, new byte[0]));

    // Create PKCS7 Signed data
    PKCS7 pkcs7 = new PKCS7(new AlgorithmId[] {digestAlgorithmId}, cInfo, new java.security.cert.X509Certificate[] {x509Certificate}, new SignerInfo[] {sInfo});

    // Write PKCS7
    ByteArrayOutputStream bOut = new DerOutputStream();
    pkcs7.encodeSignedData(bOut);
    os.write(bOut.toByteArray());
  }

  public boolean verify(PublicKey publicKey, byte[] signatureData, Consumer<NlX509Metadata> metadataConsumer, byte[] dataToVerify) throws Exception {
    boolean result = false;

    PKCS7 pkcs7 = new PKCS7(signatureData);
    SignerInfo[] signerInfos = pkcs7.getSignerInfos();

    if (signerInfos != null && signerInfos.length > 0) {
      Signature signature = Signature.getInstance(ALGORITHM);
      signature.initVerify(publicKey);
      signature.update(dataToVerify);

      result = signature.verify(signerInfos[0].getEncryptedDigest());

      X509Certificate[] x509Certificates = pkcs7.getCertificates();

      if (x509Certificates.length == 1) {
        X509Certificate x509Certificate = x509Certificates[0];

        metadataConsumer.accept(NlX509Metadata.fromDistinguihedName(x509Certificate.getSubjectDN().getName(), x509Certificate.getNotBefore(),
            x509Certificate.getNotAfter(),
            x509Certificate.getSigAlgName()));
      }
    }

    return result;
  }

  private X509Certificate getX509Certificate(PrivateKey privateKey, PublicKey publicKey, NlX509Metadata nlX509Metadata) throws Exception {
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    new NlX509().exportPublicKeyAsX509(baos, privateKey, publicKey, nlX509Metadata);

    ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
    return (X509Certificate) CertificateFactory.getInstance("X509").generateCertificate(bais);
  }
}
