package security;

import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.bouncycastle.bcpg.ArmoredOutputStream;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.openpgp.*;
import org.bouncycastle.openpgp.operator.PublicKeyDataDecryptorFactory;
import org.bouncycastle.openpgp.operator.jcajce.*;
import org.bouncycastle.util.encoders.Base64;
import org.bouncycastle.util.io.Streams;

import java.io.*;
import java.security.SecureRandom;
import java.security.SignatureException;
import java.util.Date;
import java.util.Iterator;
import java.util.List;

public class MessageLevel {

    public static boolean isBase64(String textToDecrypt) {
        try {
            Base64.decode(textToDecrypt);
            return true;
        } catch (IllegalArgumentException e) {
            return false;
        }
    }

    public static String encrypt(String textToEncrypt, PGPPublicKey bankPublicKey) {
        if (StringUtils.isBlank(textToEncrypt)) {
            return "";
        }
        ByteArrayOutputStream actualOutput = new ByteArrayOutputStream();
        try {
            ByteArrayInputStream inputStream = new ByteArrayInputStream(textToEncrypt.getBytes());
            OutputStream out = new ArmoredOutputStream(actualOutput);
            PGPEncryptedDataGenerator encGen =
                    new PGPEncryptedDataGenerator(
                            new JcePGPDataEncryptorBuilder(PGPEncryptedData.AES_256).setWithIntegrityPacket(true).setSecureRandom(
                                    new SecureRandom())
                                    .setProvider(new BouncyCastleProvider()));
            encGen.addMethod(new JcePublicKeyKeyEncryptionMethodGenerator(bankPublicKey).setProvider(new BouncyCastleProvider()));
            OutputStream encryptedOut = encGen.open(out, new byte[inputStream.available()]);
            PGPLiteralDataGenerator lGen = new PGPLiteralDataGenerator();
            OutputStream lOut = lGen.open(encryptedOut, PGPLiteralData.BINARY, "Sample-Data", new Date(),
                    new byte[inputStream.available()]); //bOut
            byte[] data = IOUtils.toByteArray(inputStream);
            lOut.write(data);
            lOut.close();
            lGen.close();
            encryptedOut.close();
            encGen.close();
            out.close();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (PGPException e) {
            e.printStackTrace();
        }
        return Base64.toBase64String(actualOutput.toByteArray());
    }

    public static String encryptAndSign(String textToEncrypt, PGPPublicKey bankPublicKey, PGPPrivateKey clientPrivateKey) {
        if (StringUtils.isBlank(textToEncrypt)) {
            return "";
        }
        ByteArrayOutputStream actualOutput = new ByteArrayOutputStream();
        try {
            ByteArrayInputStream inputStream = new ByteArrayInputStream(textToEncrypt.getBytes());
            OutputStream out = new ArmoredOutputStream(actualOutput);
            PGPEncryptedDataGenerator encGen =
                    new PGPEncryptedDataGenerator(
                            new JcePGPDataEncryptorBuilder(PGPEncryptedData.AES_256).setWithIntegrityPacket(true).setSecureRandom(
                                    new SecureRandom())
                                    .setProvider(new BouncyCastleProvider()));
            encGen.addMethod(new JcePublicKeyKeyEncryptionMethodGenerator(bankPublicKey).setProvider(new BouncyCastleProvider()));
            OutputStream encryptedOut = encGen.open(out, new byte[inputStream.available()]);
            PGPLiteralDataGenerator lGen = new PGPLiteralDataGenerator();
            PGPCompressedDataGenerator comData = new PGPCompressedDataGenerator(PGPCompressedData.ZIP);
            OutputStream compressedData = comData.open(encryptedOut);
            PGPSignatureGenerator sGen = new PGPSignatureGenerator(new JcaPGPContentSignerBuilder(
                    clientPrivateKey.getPublicKeyPacket().getAlgorithm(), PGPUtil.SHA512).setProvider(new BouncyCastleProvider()));
            sGen.init(PGPSignature.BINARY_DOCUMENT, clientPrivateKey);
            Iterator<String> it = bankPublicKey.getUserIDs();
            if (it.hasNext()) {
                PGPSignatureSubpacketGenerator spGen = new PGPSignatureSubpacketGenerator();
                spGen.setSignerUserID(false, it.next());
                sGen.setHashedSubpackets(spGen.generate());
            }
            sGen.generateOnePassVersion(false).encode(compressedData); // bOut
            OutputStream lOut = lGen.open(compressedData, PGPLiteralData.BINARY, "Sample-Data", new Date(),
                    new byte[inputStream.available()]); //bOut
            byte[] data = IOUtils.toByteArray(inputStream);
            lOut.write(data);
            sGen.update(data);

            lOut.close();
            lGen.close();
            sGen.generate().encode(compressedData);
            comData.close();
            compressedData.close();


            encryptedOut.close();
            encGen.close();
            out.close();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (PGPException e) {
            e.printStackTrace();
        }
        return Base64.toBase64String(actualOutput.toByteArray());
    }

    public static String decrypt(String textToDecrypt, List<PGPPrivateKey> clientPrivateKeys) {
        if (StringUtils.isBlank(textToDecrypt)) {
            return "";
        }
        try {
            ByteArrayOutputStream actualOutput = new ByteArrayOutputStream();
                String epicStrNew = new String(Base64.decode(textToDecrypt));
                JcaKeyFingerprintCalculator keyFingerPrintCalculator = new JcaKeyFingerprintCalculator();
                BufferedInputStream dataStream = new BufferedInputStream(IOUtils.toInputStream(epicStrNew));
                InputStream decoderStream = PGPUtil.getDecoderStream(dataStream);
                PGPObjectFactory pgpObjectFactory = new PGPObjectFactory(decoderStream, keyFingerPrintCalculator);
                PGPEncryptedDataList pgpEncryptedDataList = (PGPEncryptedDataList) pgpObjectFactory.nextObject();
                System.out.println("pgpEncryptedDataList" + pgpEncryptedDataList);
                if (pgpEncryptedDataList != null) {
                    // find the matching public key encrypted data packet.
                    PGPPublicKeyEncryptedData pgpPublicKeyEncryptedData = null;
                    // build decryptor factory
                    PublicKeyDataDecryptorFactory dataDecryptorFactory = null;
                    for (Object pgpEncData : pgpEncryptedDataList) {
                        PGPPublicKeyEncryptedData pkEnc = (PGPPublicKeyEncryptedData) pgpEncData;
                        for (PGPPrivateKey clientPrivateKey : clientPrivateKeys) {
                            if (pkEnc.getKeyID() == clientPrivateKey.getKeyID()) {
                                System.out.println(clientPrivateKey.getKeyID());
                                pgpPublicKeyEncryptedData = pkEnc;
                                dataDecryptorFactory = new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC")
                                        .build(clientPrivateKey);
                                break;
                            }
                        }
                    }
                    if (pgpPublicKeyEncryptedData == null || dataDecryptorFactory == null) {
                        // Public key used to encrypt data does not correspond to private key provided
                        throw new IllegalStateException("matching encrypted data not found");
                    }
                    InputStream clear = pgpPublicKeyEncryptedData.getDataStream(dataDecryptorFactory);
                    PGPObjectFactory plainFact = new PGPObjectFactory(clear, new JcaKeyFingerprintCalculator());
                    Object message = plainFact.nextObject();
                    // Signature related
                    PGPSignatureList signatureList = null;
                    PGPOnePassSignatureList onePassSignatureList = null;
                    PGPCompressedData compressedData = null;
                    while (message != null) {
                        if (message instanceof PGPCompressedData) {
                            compressedData = (PGPCompressedData) message;
                            plainFact = new PGPObjectFactory(compressedData.getDataStream(), keyFingerPrintCalculator);
                            message = plainFact.nextObject();
                        }
                        if (message instanceof PGPLiteralData) {
                            // have to read it and keep it somewhere.
                            Streams.pipeAll(((PGPLiteralData) message).getInputStream(), actualOutput);
                        } else if (message instanceof PGPOnePassSignatureList) {
                            onePassSignatureList = (PGPOnePassSignatureList) message;
                        } else if (message instanceof PGPSignatureList) {
                            signatureList = (PGPSignatureList) message;
                        } else {
                            throw new PGPException("message unknown message type.");
                        }
                        message = plainFact.nextObject();
                    }
                    // check data decrypts okay
                    if (!pgpPublicKeyEncryptedData.verify()) {
                        throw new RuntimeException("Enc data is invalid!");
                    }
                    actualOutput.close();
                    clear.close();
                    return actualOutput.toString();
                }
                return epicStrNew;

        } catch (PGPException e) {
            e.printStackTrace();
        } catch (IllegalArgumentException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String decryptAndVerifySignature(String textToDecrypt, List<PGPPublicKey> bankPublicKeys,
                                 List<PGPPrivateKey> clientPrivateKeys) {
        if (StringUtils.isBlank(textToDecrypt)) {
            return "";
        }
        try {
            ByteArrayOutputStream actualOutput = new ByteArrayOutputStream();
            String epicStrNew = new String(Base64.decode(textToDecrypt));
            JcaKeyFingerprintCalculator keyFingerPrintCalculator = new JcaKeyFingerprintCalculator();
            BufferedInputStream dataStream = new BufferedInputStream(IOUtils.toInputStream(epicStrNew));
            InputStream decoderStream = PGPUtil.getDecoderStream(dataStream);
            PGPObjectFactory pgpObjectFactory = new PGPObjectFactory(decoderStream, keyFingerPrintCalculator);
            PGPEncryptedDataList pgpEncryptedDataList = (PGPEncryptedDataList) pgpObjectFactory.nextObject();
            System.out.println("pgpEncryptedDataList" + pgpEncryptedDataList);
            if (pgpEncryptedDataList != null) {
                // find the matching public key encrypted data packet.
                PGPPublicKeyEncryptedData pgpPublicKeyEncryptedData = null;
                // build decryptor factory
                PublicKeyDataDecryptorFactory dataDecryptorFactory = null;
                for (Object pgpEncData : pgpEncryptedDataList) {
                    PGPPublicKeyEncryptedData pkEnc = (PGPPublicKeyEncryptedData) pgpEncData;
                    for (PGPPrivateKey clientPrivateKey : clientPrivateKeys) {
                        if (pkEnc.getKeyID() == clientPrivateKey.getKeyID()) {
                            System.out.println(clientPrivateKey.getKeyID());
                            pgpPublicKeyEncryptedData = pkEnc;
                            dataDecryptorFactory = new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC")
                                    .build(clientPrivateKey);
                            break;
                        }
                    }
                }
                if (pgpPublicKeyEncryptedData == null || dataDecryptorFactory == null) {
                    // Public key used to encrypt data does not correspond to private key provided
                    throw new IllegalStateException("matching encrypted data not found");
                }
                InputStream clear = pgpPublicKeyEncryptedData.getDataStream(dataDecryptorFactory);
                PGPObjectFactory plainFact = new PGPObjectFactory(clear, new JcaKeyFingerprintCalculator());
                Object message = plainFact.nextObject();
                // Signature related
                PGPSignatureList signatureList = null;
                PGPOnePassSignatureList onePassSignatureList = null;
                PGPCompressedData compressedData = null;
                while (message != null) {
                    if (message instanceof PGPCompressedData) {
                        compressedData = (PGPCompressedData) message;
                        plainFact = new PGPObjectFactory(compressedData.getDataStream(), keyFingerPrintCalculator);
                        message = plainFact.nextObject();
                    }
                    if (message instanceof PGPLiteralData) {
                        // have to read it and keep it somewhere.
                        Streams.pipeAll(((PGPLiteralData) message).getInputStream(), actualOutput);
                    } else if (message instanceof PGPOnePassSignatureList) {
                        onePassSignatureList = (PGPOnePassSignatureList) message;
                    } else if (message instanceof PGPSignatureList) {
                        signatureList = (PGPSignatureList) message;
                    } else {
                        throw new PGPException("message unknown message type.");
                    }
                    message = plainFact.nextObject();
                }
                // check data decrypts okay
                if (!pgpPublicKeyEncryptedData.verify()) {
                    throw new RuntimeException("Enc data is invalid!");
                }
                if (onePassSignatureList == null || signatureList == null) {
                    throw new SignatureException("PGP Signatures not found.");
                } else {
                    boolean signatureVerified = false;
                    for (int i = 0; i < onePassSignatureList.size(); i++) {
                        PGPOnePassSignature ops = onePassSignatureList.get(0);
                        PGPSignature signature = signatureList.get(0);
                        for (PGPPublicKey bankPublicKey : bankPublicKeys) {
                            System.out.println("Verification Key ID:" + Long.toHexString(bankPublicKey.getKeyID()).toUpperCase());
                            ops.init(new JcaPGPContentVerifierBuilderProvider().setProvider(new BouncyCastleProvider()),
                                    bankPublicKey);
                            ops.update(actualOutput.toByteArray());
                            if (ops.verify(signature)) {
                                signatureVerified = true;
                                break;
                            }
                        }
                    }
                    if (!signatureVerified) {
                        throw new SignatureException("Signature verification failed");
                    }
                }
                actualOutput.close();
                clear.close();
                return actualOutput.toString();
            }
            return epicStrNew;

        } catch (PGPException e) {
            e.printStackTrace();
        } catch (IllegalArgumentException e) {
            e.printStackTrace();
        } catch (SignatureException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

}
