/* 
 * This file is part of cert-encoder
 * Copyright (c) 2024 Paolo Lulli.
 * 
 * This program is free software: you can redistribute it and/or modify  
 * it under the terms of the GNU General Public License as published by  
 * the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful, but 
 * WITHOUT ANY WARRANTY; without even the implied warranty of 
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License 
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

package net.lulli.encrypt.pki;

import net.lulli.encrypt.Pems;
import org.bouncycastle.cert.jcajce.JcaCertStore;
import org.bouncycastle.cms.*;
import org.bouncycastle.cms.jcajce.JcaSignerInfoGeneratorBuilder;
import org.bouncycastle.cms.jcajce.JceCMSContentEncryptorBuilder;
import org.bouncycastle.cms.jcajce.JceKeyTransEnvelopedRecipient;
import org.bouncycastle.cms.jcajce.JceKeyTransRecipientInfoGenerator;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;
import org.bouncycastle.util.Store;
import org.bouncycastle.util.encoders.Base64;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.util.ArrayList;

public class PkiEncryptionManager {
    public static final PkiEncryptionManager INSTANCE = new PkiEncryptionManager();
    private PkiEncryptionManager() {}

    public  void signFile(File fileToSign) {
        try {
            char[] password = "yourpass".toCharArray();
            String certificateAlias = "";
            String keyAlias = "";
            var keystoreFile = "full/path/to/your/original/certificate.pfx";

            Security.addProvider(new BouncyCastleProvider());
            KeyStore ks = loadKeyStore(keystoreFile, password);

            X509Certificate cert = (X509Certificate) ks.getCertificate(certificateAlias);
            PrivateKey privatekey = (PrivateKey) ks.getKey(keyAlias, password);

            String signatureFileName = fileToSign.getName().toString() + ".signature";

            byte[] buffer = loadFileIntoBuffer(fileToSign);

            byte[] signeddata = generateSignature(cert, cert.getPublicKey(), privatekey, buffer);

            writeSignatureToFile(signatureFileName, signeddata);
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    private static byte[] loadFileIntoBuffer(File fileToSign) {
        try {
            byte[] buffer = new byte[(int) fileToSign.length()];
            try (DataInputStream in = new DataInputStream(new FileInputStream(fileToSign))) {
                in.readFully(buffer);
            }
            return buffer;
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    private static byte[] generateSignature(X509Certificate cert, PublicKey publicKey, PrivateKey privatekey, byte[] buffer) {
        try {
            ArrayList<X509Certificate> certList = new ArrayList<X509Certificate>();
            certList.add(cert);
            Store<?> certs = new JcaCertStore(certList);
            CMSSignedDataGenerator signGen = new CMSSignedDataGenerator();

            ContentSigner sha1signer = new JcaContentSignerBuilder("SHA1withRSA").setProvider("BC").build(privatekey);
            signGen.addSignerInfoGenerator(new JcaSignerInfoGeneratorBuilder(new JcaDigestCalculatorProviderBuilder().build()).build(sha1signer, cert));
            signGen.addCertificates(certs);
            CMSTypedData content = new CMSProcessableByteArray(buffer);
            CMSSignedData signedData = signGen.generate(content, false);
            byte[] signeddata = signedData.getEncoded();
            return signeddata;
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    public  byte[] encryptData(byte[] data, X509Certificate encryptionCertificate) {
        try {
            Security.addProvider(new BouncyCastleProvider());

            if (null != data && null != encryptionCertificate) {
                var cmsEnvelopedDataGenerator = new CMSEnvelopedDataGenerator();

                //JceKeyTransRecipientInfoGenerator jceKey = new JceKeyTransRecipientInfoGenerator(encryptionCertificate);
                cmsEnvelopedDataGenerator.addRecipientInfoGenerator(new JceKeyTransRecipientInfoGenerator(encryptionCertificate).setProvider("BC"));
                var cmsProcessableByteArray = new CMSProcessableByteArray(data);
                var outputEncryptor = new JceCMSContentEncryptorBuilder(CMSAlgorithm.AES256_CBC).setProvider("BC").build();
                var cmsEnvelopedData = cmsEnvelopedDataGenerator.generate(cmsProcessableByteArray, outputEncryptor);
                return cmsEnvelopedData.getEncoded();
            }
            return null;
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }


    public void decryptCms(PrivateKey privateKey, byte[] encryptedData, String destinationFile) {
        try {
            var decryptedDestination = new File(destinationFile);
            //byte[] encryptedData = Files.readAllBytes(encrypted.toPath());

            var cmsEnvelopedDataParser = new CMSEnvelopedDataParser(encryptedData);

            var singleRecipient = getSingleRecipient(cmsEnvelopedDataParser);
            var jceKeyTransEnvelopedRecipient = new JceKeyTransEnvelopedRecipient(privateKey);

            try (InputStream decryptedStream = singleRecipient.getContentStream(jceKeyTransEnvelopedRecipient).getContentStream()) {
                Files.copy(decryptedStream, decryptedDestination.toPath());
            }

            System.out.println(String.format("Decrypted  to file :'%s'", decryptedDestination.getAbsolutePath()));
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    private static RecipientInformation getSingleRecipient(CMSEnvelopedDataParser parser) {
        var recInfos = parser.getRecipientInfos().getRecipients();
        var recipientIterator = recInfos.iterator();
        if (!recipientIterator.hasNext()) {
            throw new RuntimeException("Could not find recipient");
        }
        return recipientIterator.next();
    }

    private static void writeSignatureToFile(String fileName, byte[] signeddata) {
        try {
            var fileOutputStream = new FileOutputStream(fileName);
            byte[] outputString = Base64.encode(signeddata);
            int fullLines = (int) Math.floor(outputString.length / 64);
            for (int i = 0; i < fullLines; i++) {
                fileOutputStream.write(outputString, i * 64, 64);
                fileOutputStream.write("\r\n".getBytes());
            }

            fileOutputStream.write(outputString, fullLines * 64, outputString.length % 64);
            fileOutputStream.close();
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    private static KeyStore loadKeyStore(String keystoreFile, char[] password) {
        try {
            var keyStore = KeyStore.getInstance("PKCS12");
            keyStore.load(new FileInputStream(keystoreFile), password);
            return keyStore;
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    /*
    public static void main(String[] args) throws Exception {
        var plainText = "this is a sample text";
        X509Certificate x509Certificate = Pems.certificateFromFile("/home/plulli/.x509crypt/local/certs/local.crt");
        //var enc = EncDec.encrypt(plainText, x509Certificate.getPublicKey().getEncoded());
        byte[] enc = encryptData(plainText.getBytes(), x509Certificate);
        System.out.printf("Cleartext: [%s]\n Encrypted: [%s]\n", plainText, new String(enc));

        Files.write(Paths.get("/tmp/output2.txt.enc"), enc);

        PrivateKey privateKey = Pems.privateKeyFromFile("/home/plulli/.x509crypt/local/certs/local.key");
        decryptCms(privateKey, enc, "/tmp/output2.txt");
    }

     */
}