package be.iminds.ilabt.jfed.lowlevel.ssh_key_info.putty;

import be.iminds.ilabt.jfed.util.IOUtils;
import be.iminds.ilabt.jfed.util.KeyUtil;
import be.iminds.ilabt.jfed.util.TextUtil;
import ch.ethz.ssh2.crypto.cipher.AES;
import ch.ethz.ssh2.crypto.cipher.CBCMode;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.charset.Charset;
import java.security.*;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateCrtKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * PuTTYPrivateKeyFile
 */
public class PuTTYPrivateKeyFile {
    private static class MalformedPPKException extends RuntimeException {
        private MalformedPPKException() {
        }

        private MalformedPPKException(String message) {
            super(message);
        }

        private MalformedPPKException(String message, Throwable cause) {
            super(message, cause);
        }

        private MalformedPPKException(Throwable cause) {
            super(cause);
        }
    }

    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(PuTTYPrivateKeyFile.class);


    private RSAPublicKey publicKey;
    private RSAPrivateCrtKey privateKey;
    private boolean encrypted;

    private Map<String, String> lines = new HashMap<>();

    private PuTTYPrivateKeyFile() {
    }

    public static PuTTYPrivateKeyFile parse(String content) {
        PuTTYPrivateKeyFile res = new PuTTYPrivateKeyFile();
        res.lines = contentToLines(content);
        assert Objects.equals(res.lines.get("PuTTY-User-Key-File-2"), "ssh-rsa");
        res.publicKey = linesToPublicKey(res.lines);
        res.encrypted = linesToEncrypted(res.lines);
        if (!res.isEncrypted()) res.privateKey = linesToPrivateKey(res.lines, null);
        return res;
    }

    public static PuTTYPrivateKeyFile read(File file) throws IOException {
        String content = IOUtils.fileToString(file);
        return parse(content);
    }

    public static boolean isPuttyPrivateKey(File file) throws IOException {
        BufferedReader r = new BufferedReader(new FileReader(file));
        String firstLine = r.readLine();
        r.close();
        return firstLine != null && isPuttyPrivateKey(firstLine);
    }

    public static boolean isPuttyPrivateKey(String fileContent) {
        return fileContent.startsWith("PuTTY-User-Key-File-2");
    }

    @Nonnull
    private static Map<String, String> contentToLines(@Nonnull String content) {
        Map<String, String> lines = new HashMap<>();

        try (BufferedReader r = new BufferedReader(new StringReader(content))) {
            String line;
            while ((line = r.readLine()) != null) {
                int colonPos = line.indexOf(": ");
                if (colonPos <= 0)
                    throw new MalformedPPKException("PPK contains invalid line: \"" + line + "\"");
                String key = line.substring(0, colonPos);
                String value = line.substring(colonPos + 2);
                if (!key.endsWith("-Lines")) {
                    lines.put(key, value);
                } else {
                    int lineCount = Integer.parseInt(value);
                    String fullVall = "";
                    for (int i = 0; i < lineCount; i++) {
                        String valLine = r.readLine();
                        if (valLine == null)
                            throw new MalformedPPKException("PPK ended prematurely. Was reading " + lineCount + " lines for \"" + key + "\" but got EOF at line " + i);
                        fullVall += valLine;
                    }
                    lines.put(key, fullVall);
                }
            }
        } catch (IOException ignored) {
        }

        return lines;
    }

    private static RSAPublicKey linesToPublicKey(Map<String, String> lines) {
        String pk = lines.get("Public-Lines");
        if (pk == null)
            throw new MalformedPPKException("PPK contains no \"Public-Lines\" lines");

        byte[] publicLines = Base64.decodeBase64(pk);

        ByteBuffer bb = ByteBuffer.wrap(publicLines).order(ByteOrder.BIG_ENDIAN);

        int algoLen = bb.getInt();
        byte[] algo = new byte[algoLen];
        bb.get(algo);

        int publicExponentLen = bb.getInt();
        byte[] publicExponent = new byte[publicExponentLen];
        bb.get(publicExponent);
        BigInteger publicExponentInt = new BigInteger(publicExponent);

        int modulusLen = bb.getInt();
        byte[] modulus = new byte[modulusLen];
        bb.get(modulus);
        BigInteger modulusInt = new BigInteger(modulus);

        RSAPublicKeySpec spec = new RSAPublicKeySpec(modulusInt, publicExponentInt);
        RSAPublicKey res = null;
        try {
            KeyFactory keyFact = KeyFactory.getInstance("RSA");
            PublicKey key = keyFact.generatePublic(spec);
            return (RSAPublicKey) key;
        } catch (Exception e) {
            throw new RuntimeException("Error creating RSAPublicKey: " + e.getMessage(), e);
        }
    }

    private static boolean linesToEncrypted(Map<String, String> lines) {
        String enc = lines.get("Encryption");
        if (enc == null)
            throw new MalformedPPKException("PPK contains no \"Encryption\" lines");
        if (Objects.equals(enc, "none"))
            return false;
        if (Objects.equals(enc, "aes256-cbc"))
            return true;
        throw new MalformedPPKException("PPK contains unsupported \"Encryption\" line: \"" + enc + "\"");

    }

    private static RSAPrivateCrtKey linesToPrivateKey(Map<String, String> lines, String passphrase) {
        String pk = lines.get("Private-Lines");
        if (pk == null)
            throw new MalformedPPKException("PPK contains no \"Private-Lines\" lines");

        String algorithm = lines.get("PuTTY-User-Key-File-2");
        String encryption = lines.get("Encryption");
        String publicLinesString = lines.get("Public-Lines");
        byte[] publicLines = Base64.decodeBase64(publicLinesString);

        ByteBuffer bb = ByteBuffer.wrap(publicLines).order(ByteOrder.BIG_ENDIAN);

        int algoLen = bb.getInt();
        byte[] algo = new byte[algoLen];
        bb.get(algo);

        int publicExponentLen = bb.getInt();
        byte[] publicExponent = new byte[publicExponentLen];
        bb.get(publicExponent);
        BigInteger publicExponentInt = new BigInteger(publicExponent);

        int modulusLen = bb.getInt();
        byte[] modulus = new byte[modulusLen];
        bb.get(modulus);
        BigInteger modulusInt = new BigInteger(modulus);

        String privateLinesString = lines.get("Private-Lines");
        byte[] privateLines = Base64.decodeBase64(privateLinesString);

        if (!encryption.equalsIgnoreCase("None")) {
            assert passphrase != null;
            privateLines = decryptPPkPrivateLines(encryption, passphrase, privateLines);
            //there might be some padding after the Private-Lines data, but the format ignores that.
        }

        bb = ByteBuffer.wrap(privateLines).order(ByteOrder.BIG_ENDIAN);

        int privateExponentLen = bb.getInt();
        if (privateExponentLen <= 0)
            return null; //illegal size, so decryption must have failed

        byte[] privateExponent = new byte[privateExponentLen];
        bb.get(privateExponent);
        BigInteger privateExponentInt = new BigInteger(privateExponent);

        int pLen = bb.getInt();
        if (pLen <= 0)
            return null; //illegal size, so decryption must have failed

        byte[] p = new byte[pLen];
        bb.get(p);
        BigInteger pInt = new BigInteger(p);

        int qLen = bb.getInt();
        if (qLen <= 0)
            return null; //illegal size, so decryption must have failed

        byte[] q = new byte[qLen];
        bb.get(q);
        BigInteger qInt = new BigInteger(q);

        int iqmpLen = bb.getInt();
        if (iqmpLen <= 0)
            return null; //illegal size, so decryption must have failed

        byte[] iqmp = new byte[iqmpLen];
        bb.get(iqmp);
        BigInteger iqmpInt = new BigInteger(iqmp);


        BigInteger dmp1 = privateExponentInt.mod(pInt.subtract(BigInteger.ONE));
        BigInteger dmq1 = privateExponentInt.mod(qInt.subtract(BigInteger.ONE));

        RSAPrivateCrtKeySpec rSAPrivateCrtKeySpec = new RSAPrivateCrtKeySpec(modulusInt, publicExponentInt,
                privateExponentInt, pInt, qInt, dmp1, dmq1, iqmpInt);
        KeyFactory keyFactory;
        try {
            keyFactory = KeyFactory.getInstance("RSA");
            return (RSAPrivateCrtKey) keyFactory.generatePrivate(rSAPrivateCrtKeySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            LOG.error("Error while creating Rsa Private Key: " + e.getMessage(), e);
            return null;
        }
    }

    /**
     * assumes key is unlocked!
     */
    private boolean verifyMac(String passphrase) {
        String mac = calcMac(passphrase);
        return Objects.equals(mac, lines.get("Private-MAC"));
    }


    public String getComment() {
        return lines.get("Comment");
    }

    public boolean isEncrypted() {
        return encrypted;
    }

    public boolean unlock(String passphrase) {
        privateKey = linesToPrivateKey(lines, passphrase);

        if (privateKey != null) {
            //do extra check of Mac
            if (!verifyMac(passphrase))
                throw new MalformedPPKException("PPK file MAC is incorrect.");

            return true;
        } else
            return false; //decryption has failed
    }

    public PrivateKey getPrivateKey() {
        if (privateKey != null)
            return privateKey;
        assert isEncrypted();
        throw new RuntimeException("PPK file is encrypted. Cannot retrieve private key without passphrase.");
    }

    public RSAPublicKey getPublicKey() {
        return publicKey;
    }

    public String getPublicKeyOpenSshString() {
        String algo = lines.get("PuTTY-User-Key-File-2");
        String pubkey = lines.get("Public-Lines").replace("\n", "");
        String comment = lines.get("Comment");

        return algo + " " + pubkey + " " + comment;
    }


    private static String singleLinesMaker(Map<String, String> lines, String key) {
        String val = lines.get(key);
        return key + ": " + val + "\n";
    }

    private static String multiLinesMaker(Map<String, String> lines, String key) {
        String val = lines.get(key);

        if (val.endsWith("\n")) val = val.substring(0, val.length() - 1);
        int valLineCount = (val.length() - val.replace("\n", "").length()) + 1;

        return key + ": " + valLineCount + "\n" + val + "\n";
    }

    private static String linesToContent(Map<String, String> lines) {
        return singleLinesMaker(lines, "PuTTY-User-Key-File-2") +
                singleLinesMaker(lines, "Encryption") +
                singleLinesMaker(lines, "Comment") +
                multiLinesMaker(lines, "Public-Lines") +
                multiLinesMaker(lines, "Private-Lines") +
                singleLinesMaker(lines, "Private-MAC");
    }

    private String calcMac(String passphrase) {
        try {
            MessageDigest passMd;
            passMd = MessageDigest.getInstance("SHA-1");
            passMd.update("putty-private-key-file-mac-key".getBytes(Charset.forName("US-ASCII")));
            if (passphrase != null)
                passMd.update(passphrase.getBytes(Charset.forName("US-ASCII")));
            byte[] keyBytes = passMd.digest();

            SecretKeySpec signingKey = new SecretKeySpec(keyBytes, "HmacSHA1");

            // Get an hmac_sha1 Mac instance and initialize with the signing key
            Mac mac = Mac.getInstance("HmacSHA1");
            mac.init(signingKey);

            byte[] algoBytes = lines.get("PuTTY-User-Key-File-2").getBytes(Charset.forName("US-ASCII"));
            byte[] encBytes = lines.get("Encryption").getBytes(Charset.forName("US-ASCII"));
            byte[] commentBytes = lines.get("Comment").getBytes(Charset.forName("US-ASCII"));
            byte[] pubKeyBytes = Base64.decodeBase64(lines.get("Public-Lines").getBytes(Charset.forName("US-ASCII")));
            byte[] privKeyBytes = getMacPrivateKeyBytes(privateKey, passphrase);

            int totalMacLen =
                    algoBytes.length + encBytes.length + commentBytes.length + pubKeyBytes.length + privKeyBytes.length + (5 * 4);
            byte[] macDataBytes = new byte[totalMacLen];
            ByteBuffer bbMac = ByteBuffer.wrap(macDataBytes);
            bbMac.putInt(algoBytes.length);
            bbMac.put(algoBytes);
            bbMac.putInt(encBytes.length);
            bbMac.put(encBytes);
            bbMac.putInt(commentBytes.length);
            bbMac.put(commentBytes);
            bbMac.putInt(pubKeyBytes.length);
            bbMac.put(pubKeyBytes);
            bbMac.putInt(privKeyBytes.length);
            bbMac.put(privKeyBytes);

            // Compute the hmac on input data bytes
            byte[] rawHmac = mac.doFinal(macDataBytes);

            // Convert raw bytes to Hex
            byte[] hexBytes = new Hex().encode(rawHmac);

            //  Covert array of Hex bytes to a String
            return new String(hexBytes, "UTF-8");
        } catch (NoSuchAlgorithmException | UnsupportedEncodingException | InvalidKeyException e) {
            LOG.error("Failed to calculate PPK MAC: " + e.getMessage(), e);
            return "error";
        }
    }

    private static byte[] getMacPrivateKeyBytes(RSAPrivateCrtKey rsaPrivateCrtKey, String passphrase) {
        BigInteger privExp = rsaPrivateCrtKey.getPrivateExponent();
        BigInteger p = rsaPrivateCrtKey.getPrimeP();
        BigInteger q = rsaPrivateCrtKey.getPrimeQ();
        BigInteger iqmp = rsaPrivateCrtKey.getCrtCoefficient();
        byte[] privExpBytes = privExp.toByteArray();
        byte[] pBytes = p.toByteArray();
        byte[] qBytes = q.toByteArray();
        byte[] iqmpBytes = iqmp.toByteArray();

        int totalLen = (4 * 4) + privExpBytes.length + pBytes.length + qBytes.length + iqmpBytes.length;
        byte[] privKeyBytes = new byte[totalLen];
        ByteBuffer bb = ByteBuffer.wrap(privKeyBytes).order(ByteOrder.BIG_ENDIAN);
        bb.clear();
        bb.putInt(privExpBytes.length);
        bb.put(privExpBytes);
        bb.putInt(pBytes.length);
        bb.put(pBytes);
        bb.putInt(qBytes.length);
        bb.put(qBytes);
        bb.putInt(iqmpBytes.length);
        bb.put(iqmpBytes);

        assert bb.position() == totalLen : "bb.position() (" + bb.position() + ") != totalLen (" + totalLen + ")";

        if (passphrase != null) {
            return padPPkPrivateLinesForEncryption(privKeyBytes);
        } else
            return privKeyBytes;
    }

    private static String privateKeyToLines(RSAPrivateCrtKey privateKey, String passphrase) {
        byte[] privKeyBytes = getMacPrivateKeyBytes(privateKey, passphrase);

        if (passphrase != null) {
            byte[] encyPrivKeyBytes = encryptPPkPrivateLines("aes256-cbc", passphrase, privKeyBytes);
            return TextUtil.wrap(new String(Base64.encodeBase64(encyPrivKeyBytes), Charset.forName("UTF-8")), 64);
        } else
            return TextUtil.wrap(new String(Base64.encodeBase64(privKeyBytes), Charset.forName("UTF-8")), 64);

    }

    private static String publicKeyToLines(RSAPublicKey publicKey) {
        String publicLine = KeyUtil.rsaPublicKeyToOpenSshAuthorizedKeysFormatData(publicKey);
        return TextUtil.wrap(publicLine, 64);
    }

    public static String createContent(String comment, RSAPublicKey publicKey, RSAPrivateCrtKey privateKey, String pass) {
        PuTTYPrivateKeyFile res = new PuTTYPrivateKeyFile();
        res.publicKey = publicKey;
        res.privateKey = privateKey;
        res.encrypted = pass != null;
        res.lines.put("PuTTY-User-Key-File-2", "ssh-rsa");
        res.lines.put("Comment", comment == null ? "" : comment);
        if (pass != null)
            res.lines.put("Encryption", "aes256-cbc");
        else
            res.lines.put("Encryption", "none");
        res.lines.put("Public-Lines", publicKeyToLines(publicKey));
        res.lines.put("Private-Lines", privateKeyToLines(privateKey, pass));
        res.lines.put("Private-MAC", res.calcMac(pass));
        return linesToContent(res.lines);
    }

    public static void createFile(String comment, RSAPublicKey publicKey, RSAPrivateCrtKey privateKey, File file, String pass) {
        String content = createContent(comment, publicKey, privateKey, pass);
        IOUtils.stringToFile(file, content);
    }


    /* internal stuff BELOW*/
    private static byte[] passPhraseToPuttyKey(String passphrase) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-1");

            digest.update(new byte[]{0, 0, 0, 0});
            digest.update(passphrase.getBytes());
            byte[] key1 = digest.digest();

            digest.update(new byte[]{0, 0, 0, 1});
            digest.update(passphrase.getBytes());
            byte[] key2 = digest.digest();

            byte[] r = new byte[32];
            System.arraycopy(key1, 0, r, 0, 20);
            System.arraycopy(key2, 0, r, 20, 12);

            return r;
        } catch (NoSuchAlgorithmException e) {
            throw new AssertionError(e); // impossible
        }
    }

    private static byte[] decryptPPkPrivateLines(String encryption, String pass, byte[] data) {
        assert data.length % 16 == 0 :
                "data.length=" + data.length + " data.length%16=" + (data.length % 16); //we don't pad, so this needs to be multiple of 16!
        if (Objects.equals(encryption, "aes256-cbc")) {
            AES aes = new AES();
            byte[] key = passPhraseToPuttyKey(pass);
            aes.init(false, key);
            CBCMode cbc = new CBCMode(aes, new byte[16], false); // initial List=0

            LOG.trace("DEC cbc.getBlockSize()=" + cbc.getBlockSize() + " data.length=" + data.length + " d.len/bs=" + ((data.length * 1.0) / cbc.getBlockSize()));

            byte[] res = new byte[data.length];
            for (int i = 0; i < data.length / cbc.getBlockSize(); i++) {
                cbc.transformBlock(data, i * cbc.getBlockSize(), res, i * cbc.getBlockSize());
            }
            return res;
        }
        throw new RuntimeException("Unsupported encryption: \"" + encryption + "\"");

    }

    private static byte[] padPPkPrivateLinesForEncryption(byte[] data) {
        final int BLOCK_SIZE = 16;
        int padded_size = data.length + BLOCK_SIZE - 1;
        padded_size -= padded_size % BLOCK_SIZE;
        byte[] res = new byte[padded_size];

        LOG.trace("Padding " + data.length + " bytes to " + res.length + " diff=" + (res.length - data.length) + " mods:" + (data.length % BLOCK_SIZE) + " to " + (res.length % BLOCK_SIZE));

        System.arraycopy(data, 0, res, 0, data.length);

        if (data.length % BLOCK_SIZE == 0) return res;

        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-1");
            digest.update(data);
            byte[] padding = digest.digest();

            System.arraycopy(padding, 0, res, data.length, res.length - data.length);
        } catch (NoSuchAlgorithmException e) {
            LOG.error("SHA-1 not working: " + e.getMessage(), e);
        }

        return res;
    }

    private static byte[] encryptPPkPrivateLines(String encryption, String pass, byte[] data) {
        assert data.length % 16 == 0 :
                "data.length=" + data.length + " data.length%16=" + (data.length % 16); //we don't pad, so this needs to be multiple of 16!
        if (Objects.equals(encryption, "aes256-cbc")) {
            byte[] key = passPhraseToPuttyKey(pass);
            assert key.length == 32;

            AES aes = new AES();
            aes.init(true, key);
            CBCMode cbc = new CBCMode(aes, new byte[16], true); // initial List=0

            LOG.trace("ENC cbc.getBlockSize()=" + cbc.getBlockSize() + " data.length=" + data.length + " d.len/bs=" + ((data.length * 1.0) / cbc.getBlockSize()));

            byte[] res = new byte[data.length];
            for (int i = 0; i < data.length / cbc.getBlockSize(); i++) {
                cbc.transformBlock(data, i * cbc.getBlockSize(), res, i * cbc.getBlockSize());
            }

            return res;
        }
        throw new RuntimeException("Unsupported encryption: \"" + encryption + "\"");

    }
}
