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

import com.jscape.util.ArrayTools;
import com.jscape.util.HexFormat;
import org.bouncycastle.crypto.generators.Argon2BytesGenerator;
import org.bouncycastle.crypto.params.Argon2Parameters;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.Arrays;

import static com.jscape.util.Parameters.assertNotNull;

/**
 * Putty key pair record.
 *
 * @author Alex Usun
 */
public class Record {

    private static final String NO_ENCRYPTION = "none";

    public final String algorithm;
    public final String encryptionAlgorithm;
    public final String comment;
    public final byte[] publicKeyData;
    public final byte[] privateKeyData;
    public final byte[] mac;
    public final KeyDerivation keyDerivation;

    /**
     * Creates a new record instance.
     *
     * @param algorithm           key pair algorithm
     * @param encryptionAlgorithm encryption algorithm
     * @param comment             comment
     * @param publicKeyData       public key data
     * @param privateKeyData      private key data
     * @param mac                 mac value
     */
    public Record(String algorithm,
                  String encryptionAlgorithm,
                  String comment,
                  byte[] publicKeyData,
                  byte[] privateKeyData,
                  byte[] mac,
                  KeyDerivation keyDerivation) {
        assertNotNull(algorithm);
        this.algorithm = algorithm;

        assertNotNull(encryptionAlgorithm);
        this.encryptionAlgorithm = encryptionAlgorithm;

        this.comment = comment;

        assertNotNull(publicKeyData);
        this.publicKeyData = publicKeyData;

        assertNotNull(privateKeyData);
        this.privateKeyData = privateKeyData;

        assertNotNull(mac);
        this.mac = mac;

        assertNotNull(keyDerivation);
        this.keyDerivation = keyDerivation;
    }

    /**
     * Checks if the current record is encrypted.
     *
     * @return <code>true</code> if record is encrypted; <code>false</code> otherwise
     */
    public boolean encrypted() {
        return !this.encryptionAlgorithm.equalsIgnoreCase(NO_ENCRYPTION);
    }

    @Override
    public String toString() {
        return "Record {"
                + "algorithm='" + this.algorithm + '\''
                + ", encryptionAlgorithm='" + this.encryptionAlgorithm + '\''
                + ", comment='" + this.comment + '\''
                + ", publicKeyData=" + Arrays.toString(this.publicKeyData)
                + ", privateKeyData=" + Arrays.toString(this.privateKeyData)
                + ", mac=" + Arrays.toString(this.mac)
                + ", keyDerivation=" + this.keyDerivation
                + '}';
    }

    public static class KeyData {

        public final byte[] cipherKeyData;
        public final byte[] cipherIvKeyData;
        public final byte[] macKeyData;

        public KeyData(byte[] cipherKeyData, byte[] cipherIvKeyData, byte[] macKeyData) {
            this.cipherKeyData = cipherKeyData;
            this.cipherIvKeyData = cipherIvKeyData;
            this.macKeyData = macKeyData;
        }
    }

    public interface KeyDerivation {

        KeyData keyDataFor(String password, int encryptionKeyLength, int encryptionIvLength, int macKeyLength)
                throws Exception;

    }

    public static class NullDerivation
            implements KeyDerivation {

        private static final int DIGEST_SIZE = 20;

        @Override
        public KeyData keyDataFor(String password, int encryptionKeyLength, int encryptionIvLength, int macKeyLength)
                throws Exception {
            MessageDigest digest = MessageDigest.getInstance("SHA1");
            byte[] hash = new byte[hashLengthFor(encryptionKeyLength)];
            byte[] passwordData = password.getBytes(StandardCharsets.UTF_8);
            ByteArrayDataWriter writer = new ByteArrayDataWriter();
            int counter = 0;

            for (int index = 0; index <= hash.length - DIGEST_SIZE; index += DIGEST_SIZE) {
                writer.writeInt(counter++);
                digest.update(writer.toByteArray());
                digest.update(passwordData);
                System.arraycopy(digest.digest(), 0, hash, index, DIGEST_SIZE);
            }

            return new KeyData(hash.length == encryptionKeyLength ? hash : ArrayTools.copy(hash, 0, new byte[encryptionKeyLength]),
                               new byte[encryptionIvLength],
                               null);
        }

        @Override
        public String toString() {
            return "NullDerivation {}";
        }

        private int hashLengthFor(int keySize) {
            int hashSize = Math.max(keySize, DIGEST_SIZE);
            int remainder = hashSize % DIGEST_SIZE;

            if (remainder != 0) {
                hashSize += DIGEST_SIZE - remainder;
            }

            return hashSize;
        }
    }

    public static class Argon2Derivation
            implements KeyDerivation {

        private final Argon2Parameters parameters;

        public static Argon2Derivation derivationFor(String type, String salt, int iterations, int memoryKBytes, int parallelism) {
            Argon2Parameters parameters = new Argon2Parameters.Builder(Type.typeFor(type).code)
                    .withSalt(HexFormat.parseData(salt, ""))
                    .withIterations(iterations)
                    .withMemoryAsKB(memoryKBytes)
                    .withParallelism(parallelism)
                    .build();

            return new Argon2Derivation(parameters);
        }

        public Argon2Derivation(Argon2Parameters parameters) {
            assertNotNull(parameters);
            this.parameters = parameters;
        }

        @Override
        public KeyData keyDataFor(String password, int encryptionKeyLength, int encryptionIvLength, int macKeyLength)
                throws Exception {
            byte[] hash = new byte[encryptionKeyLength + encryptionIvLength + macKeyLength];
            Argon2BytesGenerator generator = new Argon2BytesGenerator();
            generator.init(this.parameters);
            generator.generateBytes(password.getBytes(StandardCharsets.UTF_8), hash);

            return new KeyData(ArrayTools.copy(hash, 0, new byte[encryptionKeyLength]),
                               ArrayTools.copy(hash, encryptionKeyLength, new byte[encryptionIvLength]),
                               ArrayTools.copy(hash, encryptionKeyLength + encryptionIvLength, new byte[macKeyLength]));
        }

        @Override
        public String toString() {
            return "Argon2Derivation {"
                    + "parameters=" + this.parameters
                    + '}';
        }

        public enum Type {

            ARGON2_D("Argon2d", Argon2Parameters.ARGON2_d),
            ARGON2_I("Argon2i", Argon2Parameters.ARGON2_i),
            ARGON2_ID("Argon2id", Argon2Parameters.ARGON2_id);

            public final String name;
            public final int code;

            public static Type typeFor(String name) {
                for (Type candidate : values()) {
                    if (candidate.name.equals(name)) {
                        return candidate;
                    }
                }

                throw new IllegalArgumentException(String.format("Unsupported Argon2 type: %s.", name));
            }

            Type(String name, int code) {
                this.name = name;
                this.code = code;
            }
        }
    }
}
