/*
 * 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 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.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.security.KeyStore;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.util.ArrayList;

public class PkiManager {
    public static final PkiManager INSTANCE = new PkiManager();

    private PkiManager() {
    }

    private 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 = sign(cert, 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());
        }
    }

    public byte[] sign(X509Certificate cert, PrivateKey privatekey, byte[] buffer) {
        try {
            Security.addProvider(new BouncyCastleProvider());
            ArrayList<X509Certificate> certList = new ArrayList<X509Certificate>();
            certList.add(cert);
            Store<?> certs = new JcaCertStore(certList);
            var cmsSignedDataGenerator = new CMSSignedDataGenerator();

            var sha1signer = new JcaContentSignerBuilder("SHA1withRSA").setProvider("BC").build(privatekey);
            cmsSignedDataGenerator.addSignerInfoGenerator(new JcaSignerInfoGeneratorBuilder(new JcaDigestCalculatorProviderBuilder().build()).build(sha1signer, cert));
            cmsSignedDataGenerator.addCertificates(certs);
            var cmsProcessableByteArray = new CMSProcessableByteArray(buffer);
            var cmsSignedData = cmsSignedDataGenerator.generate(cmsProcessableByteArray, false);
            return cmsSignedData.getEncoded();
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

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

            if ((null == data) || (null == x509Certificate)) {
                throw new IllegalStateException("Null parameters");
            }
            var cmsEnvelopedDataGenerator = new CMSEnvelopedDataGenerator();
            cmsEnvelopedDataGenerator.addRecipientInfoGenerator(new JceKeyTransRecipientInfoGenerator(x509Certificate).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();
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    public byte[] decrypt(PrivateKey privateKey, byte[] encryptedData) {
        try {
            var decryptedStream = decryptToInputStream(privateKey, encryptedData);
            return decryptedStream.readAllBytes();
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    public void decryptToFile(PrivateKey privateKey, byte[] encryptedData, String destinationFile) {
        try {
            var decryptedDestination = new File(destinationFile);
            var decryptedStream = decryptToInputStream(privateKey, encryptedData);

            Files.copy(decryptedStream, decryptedDestination.toPath());
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage());
        }
    }

    private InputStream decryptToInputStream(PrivateKey privateKey, byte[] encryptedData) {
        try {
            Security.addProvider(new BouncyCastleProvider());
            var cmsEnvelopedDataParser = new CMSEnvelopedDataParser(encryptedData);

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

            return singleRecipient.getContentStream(jceKeyTransEnvelopedRecipient).getContentStream();
        } 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());
        }
    }
}