package info.knacki.pass.io.pgp;

import android.content.Context;
import android.support.annotation.NonNull;

import org.bouncycastle.bcpg.ArmoredOutputStream;
import org.bouncycastle.bcpg.CompressionAlgorithmTags;
import org.bouncycastle.bcpg.HashAlgorithmTags;
import org.bouncycastle.bcpg.sig.KeyFlags;
import org.bouncycastle.crypto.generators.RSAKeyPairGenerator;
import org.bouncycastle.crypto.params.RSAKeyGenerationParameters;
import org.bouncycastle.openpgp.PGPCompressedData;
import org.bouncycastle.openpgp.PGPCompressedDataGenerator;
import org.bouncycastle.openpgp.PGPEncryptedData;
import org.bouncycastle.openpgp.PGPEncryptedDataGenerator;
import org.bouncycastle.openpgp.PGPEncryptedDataList;
import org.bouncycastle.openpgp.PGPException;
import org.bouncycastle.openpgp.PGPKeyPair;
import org.bouncycastle.openpgp.PGPKeyRingGenerator;
import org.bouncycastle.openpgp.PGPLiteralData;
import org.bouncycastle.openpgp.PGPLiteralDataGenerator;
import org.bouncycastle.openpgp.PGPObjectFactory;
import org.bouncycastle.openpgp.PGPPBEEncryptedData;
import org.bouncycastle.openpgp.PGPPrivateKey;
import org.bouncycastle.openpgp.PGPPublicKey;
import org.bouncycastle.openpgp.PGPPublicKeyEncryptedData;
import org.bouncycastle.openpgp.PGPSecretKey;
import org.bouncycastle.openpgp.PGPSecretKeyRing;
import org.bouncycastle.openpgp.PGPSecretKeyRingCollection;
import org.bouncycastle.openpgp.PGPSignature;
import org.bouncycastle.openpgp.PGPSignatureSubpacketGenerator;
import org.bouncycastle.openpgp.PGPUtil;
import org.bouncycastle.openpgp.operator.PBESecretKeyDecryptor;
import org.bouncycastle.openpgp.operator.PBESecretKeyEncryptor;
import org.bouncycastle.openpgp.operator.bc.BcPBEDataDecryptorFactory;
import org.bouncycastle.openpgp.operator.bc.BcPGPContentSignerBuilder;
import org.bouncycastle.openpgp.operator.bc.BcPGPDigestCalculatorProvider;
import org.bouncycastle.openpgp.operator.bc.BcPGPKeyPair;
import org.bouncycastle.openpgp.operator.bc.BcPublicKeyDataDecryptorFactory;
import org.bouncycastle.openpgp.operator.jcajce.JcaKeyFingerprintCalculator;
import org.bouncycastle.openpgp.operator.jcajce.JcaPGPDigestCalculatorProviderBuilder;
import org.bouncycastle.openpgp.operator.jcajce.JcePBESecretKeyDecryptorBuilder;
import org.bouncycastle.openpgp.operator.jcajce.JcePBESecretKeyEncryptorBuilder;
import org.bouncycastle.openpgp.operator.jcajce.JcePGPDataEncryptorBuilder;
import org.bouncycastle.openpgp.operator.jcajce.JcePublicKeyKeyEncryptionMethodGenerator;
import org.bouncycastle.util.io.Streams;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Date;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.logging.Level;
import java.util.logging.Logger;

import info.knacki.pass.io.FileInterfaceFactory;
import info.knacki.gitdroid.callback.OnResponseListener;

public class GPGUtil {
    private final static Logger log = Logger.getLogger(GPGUtil.class.getName());

    public static class MalformedKeyException extends IOException {
        final Throwable fCause;

        MalformedKeyException(Throwable e) {
            fCause = e;
        }

        @Override
        public synchronized Throwable getCause() {
            return fCause;
        }

        @Override
        public String getMessage() {
            return fCause.getMessage();
        }
    }

    private static class PGPPrivateKeyAndPass {
        final PGPPrivateKey fPrivateKey;
        final String fPass;

        PGPPrivateKeyAndPass(PGPPrivateKey privateKey, String key) {
            fPrivateKey = privateKey;
            fPass = key;
        }
    }

    private static PBESecretKeyDecryptor CreateDecryptor(char[] pass) throws PGPException {
        return new JcePBESecretKeyDecryptorBuilder(new JcaPGPDigestCalculatorProviderBuilder().setProvider("BC").build()).setProvider("BC").build(pass);
    }

    private static PGPPrivateKeyAndPass TryPassword(PGPSecretKey pgpSecKey, String pass) {
        try {
            return TryPasswordUnsafe(pgpSecKey, pass);
        } catch (PGPException e) {
            // Wrong password
            log.log(Level.INFO, e.getMessage() + " (wrong password ?)", e);
            return null;
        }
    }

    private static @NonNull
    PGPPrivateKeyAndPass TryPasswordUnsafe(PGPSecretKey pgpSecKey, String pass) throws PGPException {
        return new PGPPrivateKeyAndPass(pgpSecKey.extractPrivateKey(CreateDecryptor(pass.toCharArray())), pass);
    }

    private static void FindPassword(final FileInterfaceFactory.PasswordGetter passwordGetter, final PGPSecretKey pgpSecKey, final OnResponseListener<PGPPrivateKeyAndPass> resp) {
        PGPPrivateKeyAndPass privateKey = TryPassword(pgpSecKey, "");

        if (privateKey != null) {
            resp.OnResponse(privateKey);
        } else {
            FileInterfaceFactory.OnPasswordEnteredListener onResp = new FileInterfaceFactory.OnPasswordEnteredListener() {
                @Override
                public boolean OnResponse(String result) {
                    if (result == null) {
                        resp.OnError("Invalid password", null);
                        return false;
                    } else {
                        PGPPrivateKeyAndPass privateKey = TryPassword(pgpSecKey, result);
                        if (privateKey != null) {
                            resp.OnResponse(privateKey);
                            return true;
                        } else {
                            passwordGetter.GetPassword(this);
                            return false;
                        }
                    }
                }

                @Override
                public void OnError(String msg, Throwable e) {
                    resp.OnError(msg, e);
                }
            };
            passwordGetter.GetPassword(onResp);
        }
    }

    private static InputStream getFirstLiteralDataInputStream(InputStream in, PGPPrivateKeyAndPass pKey) throws IOException, PGPException {
        PGPObjectFactory pof = new PGPObjectFactory(in, new JcaKeyFingerprintCalculator());
        Object po;

        while (true) {
            po = pof.nextObject();

            if (po == null) {
                break;
            } else if (po instanceof PGPCompressedData) {
                InputStream ret = getFirstLiteralDataInputStream(((PGPCompressedData) po).getDataStream(), pKey);
                if (ret != null)
                    return ret;
            } else if (po instanceof PGPLiteralData) {
                return ((PGPLiteralData) po).getDataStream();
            } else if (po instanceof PGPEncryptedDataList) {
                for (Iterator it = ((PGPEncryptedDataList) po).getEncryptedDataObjects(); it.hasNext(); ) {
                    PGPEncryptedData ped = (PGPEncryptedData) it.next();

                    if (ped instanceof PGPPublicKeyEncryptedData) {
                        InputStream ret = getFirstLiteralDataInputStream(((PGPPublicKeyEncryptedData) ped).getDataStream(new BcPublicKeyDataDecryptorFactory(pKey.fPrivateKey)), pKey);

                        if (ret != null)
                            return ret;
                    } else if (ped instanceof PGPPBEEncryptedData) {
                        InputStream ret = getFirstLiteralDataInputStream(((PGPPBEEncryptedData) ped).getDataStream(new BcPBEDataDecryptorFactory(pKey.fPass.toCharArray(), new BcPGPDigestCalculatorProvider())), pKey);
                        if (ret != null)
                            return ret;
                    } else {
                        throw new IOException("Unknown encryption packet");
                    }
                }
            }
        }
        return null;
    }

    private static void DoDecryptFile(PGPPrivateKeyAndPass sKey, InputStream dataStream, OnResponseListener<byte[]> resp) {
        byte[] output;
        try {
            InputStream ss = getFirstLiteralDataInputStream(dataStream, sKey);
            if (null == ss) {
                resp.OnError("No literal GPG data found", null);
                return;
            }
            ByteArrayOutputStream actualOutput = new ByteArrayOutputStream();
            Streams.pipeAll(ss, actualOutput);
            actualOutput.close();
            output = actualOutput.toByteArray();
        } catch (PGPException | IOException e) {
            resp.OnError(e.getMessage(), e);
            return;
        }
        resp.OnResponse(output);
    }

    private static PGPSecretKeyRing DoChangePassword(PGPSecretKeyRing secKeyring, String oldPassword, String newPassword) throws PGPException {
        final PBESecretKeyDecryptor keyDecryptor = CreateDecryptor(oldPassword.toCharArray());
        final PBESecretKeyEncryptor newEncryptor = new JcePBESecretKeyEncryptorBuilder(PGPEncryptedData.AES_256).setProvider("BC").build(newPassword.toCharArray());
        return PGPSecretKeyRing.copyWithNewPassword(secKeyring, keyDecryptor, newEncryptor);
    }

    public static void DecryptFile(final Context ctx, final FileInterfaceFactory.PasswordGetter passwordGetter, final File file, final OnResponseListener<byte[]> resp) throws IOException {
        PGPObjectFactory pgpF = new PGPObjectFactory(PGPUtil.getDecoderStream(new FileInputStream(file)), new JcaKeyFingerprintCalculator());
        Object o = pgpF.nextObject();
        final Iterator<?> it = ((o instanceof PGPEncryptedDataList) ? (PGPEncryptedDataList) o : (PGPEncryptedDataList) pgpF.nextObject()).getEncryptedDataObjects();

        while (it.hasNext()) {
            final PGPPublicKeyEncryptedData pbe = (PGPPublicKeyEncryptedData) it.next();

            FindPassword(
                    passwordGetter,
                    GPGStorageEngine.GetDefaultEngine(ctx).GetGPGSecretKey(pbe.getKeyID()),
                    new OnResponseListener<PGPPrivateKeyAndPass>() {
                        @Override
                        public void OnResponse(PGPPrivateKeyAndPass sKey) {
                            try {
                                DoDecryptFile(sKey, pbe.getDataStream(new BcPublicKeyDataDecryptorFactory(sKey.fPrivateKey)), resp);
                            } catch (PGPException e) {
                                OnError(e.getMessage(), e);
                            }
                        }

                        @Override
                        public void OnError(String msg, Throwable e) {
                            resp.OnError(msg, e);
                        }
                    });
        }
    }

    public static void CryptFile(final Context ctx, final OutputStream fileOutStream, final byte[] data, final OnResponseListener<Void> resp) {
        try {
            ByteArrayOutputStream compressedDataStream = new ByteArrayOutputStream();
            PGPCompressedDataGenerator comData = new PGPCompressedDataGenerator(CompressionAlgorithmTags.UNCOMPRESSED);
            OutputStream cos = comData.open(compressedDataStream);
            PGPLiteralDataGenerator lData = new PGPLiteralDataGenerator();

            OutputStream pOut = lData.open(cos, // the compressed output stream
                    PGPLiteralData.BINARY,
                    PGPLiteralData.CONSOLE,  // "filename" to store
                    data.length, // length of clear data
                    new Date()  // current time
            );

            pOut.write(data);
            pOut.close();
            comData.close();

            final byte[] compressedData = compressedDataStream.toByteArray();
            final OutputStream out = new ArmoredOutputStream(fileOutStream);
            final PGPSecretKey sKey = GPGStorageEngine.GetDefaultEngine(ctx).GetGPGSecretKey();

            PGPEncryptedDataGenerator encGen = new PGPEncryptedDataGenerator(new JcePGPDataEncryptorBuilder(PGPEncryptedDataGenerator.AES_128).setWithIntegrityPacket(true).setSecureRandom(new SecureRandom()).setProvider("BC"));
            encGen.addMethod(new JcePublicKeyKeyEncryptionMethodGenerator(sKey.getPublicKey()).setProvider("BC"));
            try {
                OutputStream encOut = encGen.open(fileOutStream, compressedData.length);

                encOut.write(compressedData);
                encOut.close();
                out.close();
                fileOutStream.close();
                resp.OnResponse(null);
            } catch (IOException | PGPException e) {
                resp.OnError(e.getMessage(), e);
            }
        } catch (IOException e) {
            resp.OnError(e.getMessage(), e);
        }
    }

    public static void ChangePassword(final Context ctx, final FileInterfaceFactory.ChangePasswordGetter passwordGetter, final OnResponseListener<Void> onDone) throws IOException {
        final PGPSecretKeyRing secretKeyring = GPGStorageEngine.GetDefaultEngine(ctx).GetGPGKeyRing();
        final PGPSecretKey signingKey = secretKeyring.getSecretKey();

        FindPassword(
                passwordGetter.SetStep(FileInterfaceFactory.ChangePasswordGetter.eStep.eOldPassword),
                signingKey,
                new OnResponseListener<PGPPrivateKeyAndPass>() {
                    @Override
                    public void OnResponse(final PGPPrivateKeyAndPass sKey) {
                        passwordGetter.SetStep(FileInterfaceFactory.ChangePasswordGetter.eStep.eNewPassword).GetPassword(new FileInterfaceFactory.OnPasswordEnteredListener() {
                            @Override
                            public boolean OnResponse(String result) {
                                if (result != null) {
                                    try {
                                        ArrayList<PGPSecretKeyRing> keyringCollection = new ArrayList<>();
                                        keyringCollection.add(DoChangePassword(secretKeyring, sKey.fPass, result));
                                        ByteArrayOutputStream stream = new ByteArrayOutputStream();
                                        new PGPSecretKeyRingCollection(keyringCollection).encode(stream);
                                        GPGStorageEngine.GetDefaultEngine(ctx).SetGPGKeyContent(new ByteArrayInputStream(stream.toByteArray()));
                                        onDone.OnResponse(null);
                                        return true;
                                    } catch (PGPException | IOException e) {
                                        OnError(e.getMessage(), e);
                                    }
                                }
                                return false;
                            }

                            @Override
                            public void OnError(String msg, Throwable e) {
                                onDone.OnError(msg, e);
                            }
                        });
                    }

                    @Override
                    public void OnError(String msg, Throwable e) {
                    }
                });
    }

    public static String GetGPGKeyName(Context ctx) {
        try {
            PGPSecretKey key = GPGStorageEngine.GetDefaultEngine(ctx).GetGPGSecretKey();
            Iterator<String> users = key.getUserIDs();
            if (users.hasNext())
                return users.next();
            return "" +key.getKeyID();
        } catch (IOException | NoSuchElementException e) {
            log.log(Level.WARNING, e.getMessage(), e);
        }
        return "";
    }

    public static byte[] Generate(String user) {
        try {
            final Date now = new Date();
            RSAKeyPairGenerator rsaKeygen = new RSAKeyPairGenerator();
            rsaKeygen.init(new RSAKeyGenerationParameters(BigInteger.valueOf(257), new SecureRandom(), 2048, 12));
            PGPKeyPair rsaEncryptor = new BcPGPKeyPair(PGPPublicKey.RSA_GENERAL, rsaKeygen.generateKeyPair(), now);

            PGPSignatureSubpacketGenerator encHashGen = new PGPSignatureSubpacketGenerator();
            encHashGen.setKeyFlags(false, KeyFlags.ENCRYPT_COMMS | KeyFlags.ENCRYPT_STORAGE);

            final PBESecretKeyEncryptor newEncryptor = new JcePBESecretKeyEncryptorBuilder(PGPEncryptedData.AES_256).setProvider("BC").build("".toCharArray());
            PGPKeyRingGenerator keygen = new PGPKeyRingGenerator(
                    PGPSignature.POSITIVE_CERTIFICATION,
                    rsaEncryptor,
                    user,
                    new BcPGPDigestCalculatorProvider().get(HashAlgorithmTags.SHA1),
                    encHashGen.generate(),
                    null,
                    new BcPGPContentSignerBuilder(
                            rsaEncryptor.getPublicKey().getAlgorithm(),
                            HashAlgorithmTags.SHA1),
                    newEncryptor);
            ArrayList<PGPSecretKeyRing> keyringCollection = new ArrayList<>();
            keyringCollection.add(keygen.generateSecretKeyRing());
            ByteArrayOutputStream stream = new ByteArrayOutputStream();
            new PGPSecretKeyRingCollection(keyringCollection).encode(stream);
            return stream.toByteArray();
        }
        catch (IOException | PGPException e) {
            log.log(Level.SEVERE, "Cannot generate new key", e);
            return null;
        }
    }

    public static boolean CheckIsPasswordProtected(Context ctx) throws IOException {
        return TryPassword(GPGStorageEngine.GetDefaultEngine(ctx).GetGPGSecretKey(), "") == null;
    }
}
