/*
 * Class: PuttyFormat
 *
 * Author: Alex Usun
 *
 * Date: 1/9/2005
 */
package com.jscape.inet.ssh.util.keyreader.putty;

import com.jscape.inet.ssh.util.keyreader.FormatException;
import com.jscape.inet.ssh.util.keyreader.KeyPairFormat;
import com.jscape.util.Assert;
import com.jscape.util.Base64Format;
import com.jscape.util.HexFormat;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.StringReader;
import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.KeyPair;
import java.util.HashMap;
import java.util.Map;
import java.util.Scanner;

/**
 * Putty key pair format. Here is format description taken from Putty sources:
 * <p/>
 * PuTTY's own format for SSH-2 keys is as follows:
 * <p/>
 * The file is text. Lines are terminated by CRLF, although CR-only and LF-only are tolerated on input.
 * <p/>
 * The first line says "PuTTY-User-Key-File-2: " plus the name of the algorithm ("ssh-dss", "ssh-rsa" etc).
 * <p/>
 * The next line says "Encryption: " plus an encryption type. Currently the only supported encryption types are
 * "aes256-cbc" and "none".
 * <p/>
 * The next line says "Comment: " plus the comment string.
 * <p/>
 * Next there is a line saying "Public-Lines: " plus a number N. The following N lines contain a base64 encoding of the
 * public part of the key. This is encoded as the standard SSH-2 public key blob (with no initial length): so for RSA,
 * for example, it will read
 * <p/>
 * string "ssh-rsa"
 * mpint  exponent
 * mpint  modulus
 * <p/>
 * Next, there is a line saying "Private-Lines: " plus a number N, and then N lines containing the (potentially
 * encrypted) private part of the key. For the key type "ssh-rsa", this will be composed of
 * <p/>
 * mpint  private_exponent
 * mpint  p (the larger of the two primes)
 * mpint  q (the smaller prime)
 * mpint  iqmp (the inverse of q modulo p)
 * data   padding (to reach a multiple of the cipher block size)
 * <p/>
 * And for "ssh-dss", it will be composed of
 * <p/>
 * mpint  x (the private key parameter)
 * [ string hash   20-byte hash of mpints p || q || g   only in old format ]
 * <p/>
 * Finally, there is a line saying "Private-MAC: " plus a hex representation of a HMAC-SHA-1 of:
 * <p/>
 * string  name of algorithm ("ssh-dss", "ssh-rsa")
 * string  encryption type
 * string  comment
 * string  public-blob
 * string  private-plaintext (the plaintext version of the private part, including the final padding)
 * <p/>
 * The key to the MAC is itself a SHA-1 hash of:
 * <p/>
 * data    "putty-private-key-file-mac-key"
 * data    passphrase
 * <p/>
 * (An empty passphrase is used for unencrypted keys.)
 * <p/>
 * If the key is encrypted, the encryption key is derived from the passphrase by means of a succession of SHA-1 hashes.
 * Each hash is the hash of:
 * <p/>
 * uint32  sequence-number
 * data    passphrase
 * <p/>
 * where the sequence-number increases from zero. As many of these hashes are used as necessary.
 * <p/>
 * For backwards compatibility with snapshots between 0.51 and 0.52, the older key file format, which begins with
 * "PuTTY-User-Key-File-1" (version number differs) is also supported. In this format the Private-MAC: field only covers
 * the private-plaintext field and nothing else (and without the 4-byte string length on the front too). Moreover, the
 * Private-MAC: field can be replaced with a Private-Hash: field which is a plain SHA-1 hash instead of an HMAC (this
 * was generated for unencrypted keys).
 *
 * @author Alex Usun
 * @version 1.0
 * @see <a href="PuttyFormat.java.html">Source code</a>
 */
public abstract class PuttyFormat
        implements KeyPairFormat {

    /**
     * @see KeyPairFormat#restoreKeyPair(byte[], String)
     */
    public KeyPair restoreKeyPair(byte[] data, String passphrase)
            throws FormatException {
        Assert.notNull(data);
        Assert.notNull(passphrase);


        try {
            Record record = recordFor(data);
            byte[] privateKeyData = record.encrypted() ? decryptPrivateKey(record, passphrase) : record.privateKeyData;

            return keyPairFor(record.publicKeyData, privateKeyData);
        } catch (Exception e) {
            throw new FormatException(e);
        }
    }

    protected abstract KeyPair keyPairFor(byte[] publicKeyData, byte[] privateKeyData)
            throws Exception;

    /**
     * Returns supported key pair algorithm name.
     *
     * @return supported key pair algorithm name
     */
    protected abstract String sshAlgorithm();

    /**
     * Restores key pair record.
     *
     * @param data record data
     *
     * @return restored key pair record
     *
     * @throws IllegalArgumentException if record cannot be restored
     */
    private Record recordFor(byte[] data) {
        try {
            Map<String, String> elements = parse(new String(data, StandardCharsets.US_ASCII));

            String keyAlgorithm = elements.get(elements.keySet().stream().filter(k -> k.startsWith("PuTTY-User-Key-File-")).findFirst().get());
            Assert.isTrue(sshAlgorithm().equalsIgnoreCase(keyAlgorithm), String.format("Unsupported algorithm: %s.", keyAlgorithm));

            String encryptionAlgorithm = elements.get("Encryption");
            String comment = elements.get("Comment");

            byte[] publicKeyData = Base64Format.parseData(elements.get("Public-Lines"));
            byte[] privateKeyData = Base64Format.parseData(elements.get("Private-Lines"));
            byte[] mac = HexFormat.parseData(elements.get("Private-MAC"), "");

            Record.KeyDerivation keyDerivation = keyDerivationOf(elements);

            return new Record(keyAlgorithm, encryptionAlgorithm, comment, publicKeyData, privateKeyData, mac, keyDerivation);
        } catch (Exception e) {
            throw new IllegalArgumentException("Unsupported key format.");
        }
    }

    private Map<String, String> parse(String data)
            throws Exception {
        BufferedReader reader = new BufferedReader(new StringReader(data));
        Map<String, String> elements = new HashMap<>();
        String line;
        while ((line = reader.readLine()) != null) {
            Scanner scanner = new Scanner(line).useDelimiter(":");
            String name = scanner.next().trim();
            String value = scanner.hasNext() ? scanner.skip(":").nextLine().trim() : "";

            switch (name) {
                case "Public-Lines":
                case "Private-Lines":
                    int lines = Integer.parseInt(value);
                    value = readLines(reader, lines);
                    break;
            }

            elements.put(name, value);
        }

        return elements;
    }

    /**
     * Reads specified count lines.
     *
     * @param reader line reader
     * @param count  target lines count
     *
     * @return readed lines as one string
     *
     * @throws IOException if an I/O error occurs
     */
    private String readLines(BufferedReader reader, int count)
            throws IOException {
        StringBuilder buffer = new StringBuilder();

        for (int i = 0; i < count; ++i) {
            String line = reader.readLine();

            if (line != null) {
                buffer.append(line);
            } else {
                throw new IOException("Unexpected EOF found.");
            }
        }

        return buffer.toString();
    }

    private Record.KeyDerivation keyDerivationOf(Map<String, String> elements) {
        String derivationType = elements.get("Key-Derivation");

        if (derivationType == null) {
            return new Record.NullDerivation();
        }

        if (derivationType.startsWith("Argon2")) {
            String salt = elements.get("Argon2-Salt");
            int iterations = Integer.parseInt(elements.get("Argon2-Passes"));
            int memoryKBytes = Integer.parseInt(elements.get("Argon2-Memory"));
            int parallelism = Integer.parseInt(elements.get("Argon2-Parallelism"));

            return Record.Argon2Derivation.derivationFor(derivationType, salt, iterations, memoryKBytes, parallelism);
        }

        throw new IllegalArgumentException(String.format("Unsupported key derivation: %s.", derivationType));
    }

    private byte[] decryptPrivateKey(Record record, String password)
            throws Exception {
        Cipher cipher = Cipher.getInstance("AES/CBC/NoPadding");
        int keyLengthBytes = 32;
        Record.KeyData keyData = record.keyDerivation.keyDataFor(password, keyLengthBytes, cipher.getBlockSize(), 32);
        Key key = new SecretKeySpec(keyData.cipherKeyData, 0, keyLengthBytes, "AES");
        cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(keyData.cipherIvKeyData));

        return cipher.doFinal(record.privateKeyData);
    }

}
