package com.jscape.inet.ssh.protocol.v2.transport.keyexchange;

import com.jscape.crypto.Hash;
import com.jscape.crypto.JceHash;
import com.jscape.inet.ssh.protocol.messages.IdentificationString;
import com.jscape.inet.ssh.protocol.v2.marshaling.MessageCodec;
import com.jscape.inet.ssh.protocol.v2.marshaling.keys.EcPublicKeyCodec;
import com.jscape.inet.ssh.protocol.v2.marshaling.keys.EcSignatures;
import com.jscape.inet.ssh.protocol.v2.marshaling.keys.PublicKeyCodec;
import com.jscape.inet.ssh.protocol.v2.marshaling.keys.PublicKeyCodecs;
import com.jscape.inet.ssh.protocol.v2.marshaling.primitive.MpintCodec;
import com.jscape.inet.ssh.protocol.v2.marshaling.primitive.StringCodec;
import com.jscape.inet.ssh.protocol.v2.messages.NamedEllipticCurveDomainParameters;
import com.jscape.inet.ssh.protocol.v2.messages.SshMsgKexInit;
import com.jscape.util.ArrayTools;
import com.jscape.util.io.ByteArrayOutputStream;
import com.jscape.util.io.DataCodec;

import javax.crypto.KeyAgreement;
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.Provider;
import java.security.PublicKey;
import java.security.interfaces.ECPublicKey;
import java.security.spec.*;
import java.util.Map;

import static com.jscape.util.CollectionTools.map;
import static com.jscape.util.Pair.pair;
import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Uzun
 */
public class EcDiffieHellmanKeyExchange {

    protected static final String KEY_ALGORITHM = "ECDH";

    protected static final int MAX_KEY_GENERATION_ATTEMPTS = 1000;

    protected static final Map<String, Integer> ACCEPTABLE_COORDINATE_SIZE = map(
            pair("1.2.840.10045.3.1.7", 32),
            pair("1.3.132.0.34", 48),
            pair("1.3.132.0.35", 66)
    );

    protected final String curveIdentifier;
    protected final String curveOid;
    protected final Object provider;
    protected final MessageCodec.Factory codecFactory;

    public EcDiffieHellmanKeyExchange(String curveIdentifier, Object provider) {
        assertNotNull(curveIdentifier);
        this.curveIdentifier = curveIdentifier;

        this.curveOid = NamedEllipticCurveDomainParameters.oidFor(this.curveIdentifier);

        this.provider = provider;
        this.codecFactory = EcDiffieHellmanKeyExchangeMessages.defaultMessages(this.curveIdentifier, provider);
    }

    protected KeyPair generateKeyPair()
            throws Exception {
        KeyPairGenerator generator = this.provider instanceof String ? KeyPairGenerator.getInstance(KEY_ALGORITHM, (String) this.provider) : this.provider instanceof Provider ? KeyPairGenerator.getInstance(KEY_ALGORITHM, (Provider) this.provider) : KeyPairGenerator.getInstance(KEY_ALGORITHM);
        generator.initialize(new ECGenParameterSpec(this.curveOid));

        for (int i = 0; i < MAX_KEY_GENERATION_ATTEMPTS; ++i) {
            KeyPair candidate = generator.generateKeyPair();

            if (keyPairValid(candidate)) {
                return candidate;
            }
        }

        throw new Exception("Error generation valid DH key pair.");
    }

    private boolean keyPairValid(KeyPair keyPair) {
        ECPublicKey publicKey = (ECPublicKey) keyPair.getPublic();
        ECPoint w = publicKey.getW();

        if (w.equals(ECPoint.POINT_INFINITY)) {
            return false;
        }

        ECParameterSpec params = publicKey.getParams();
        EllipticCurve curve = params.getCurve();
        BigInteger p = ((ECFieldFp) curve.getField()).getP();
        BigInteger upperBound = p.subtract(BigInteger.ONE);

        byte[] r = removeLeadingZeroes(w.getAffineX().toByteArray());
        BigInteger affineX = new BigInteger(1, r);

        byte[] s = removeLeadingZeroes(w.getAffineY().toByteArray());
        BigInteger affineY = new BigInteger(1, s);

        if (r.length != s.length || r.length != ACCEPTABLE_COORDINATE_SIZE.getOrDefault(this.curveOid, r.length)) {
            return false;
        }

        if (!(affineX.compareTo(upperBound) <= 0 && affineY.compareTo(upperBound) <= 0)) {
            return false;
        }

        // y^2 = x^3 + x*a + b (mod p)
        return affineY.modPow(BigInteger.ONE.add(BigInteger.ONE), p).equals(affineX.multiply(curve.getA()).add(curve.getB()).add(affineX.modPow(BigInteger.ONE.add(BigInteger.ONE).add(BigInteger.ONE), p)).mod(p));
    }

    protected byte[] removeLeadingZeroes(byte[] data) {
        int i = 0;
        while (data[i] == 0x00) {
            ++i;
        }

        return i == 0 ? data : ArrayTools.copy(data, i, new byte[data.length - i]);
    }

    protected JceHash hashFor(PublicKey key)
            throws Exception {
        String hashAlgorithm = EcSignatures.hashAlgorithmFor(key);

        return this.provider instanceof String
               ? JceHash.hashFor(hashAlgorithm, (String) this.provider)
               : this.provider instanceof Provider ? JceHash.hashFor(hashAlgorithm, (Provider) this.provider)
                                                   : JceHash.hashFor(hashAlgorithm, (String) null);
    }

    protected byte[] computeSharedSecret(KeyPair localKeys, PublicKey remoteKey)
            throws Exception {
        KeyAgreement keyAgreement = this.provider instanceof String
                                    ? KeyAgreement.getInstance("ECDH", (String) this.provider)
                                    : this.provider instanceof Provider ? KeyAgreement.getInstance("ECDH", (Provider) this.provider)
                                                                        : KeyAgreement.getInstance("ECDH");
        keyAgreement.init(localKeys.getPrivate());
        keyAgreement.doPhase(remoteKey, true);
        byte[] sharedSecretData = removeLeadingZeroes(keyAgreement.generateSecret());
        BigInteger sharedSecret = new BigInteger(1, sharedSecretData);

        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        MpintCodec.writeValue(sharedSecret, buffer);

        return buffer.toByteArray();
    }

    protected byte[] exchangeHashFor(Hash hash,
                                     IdentificationString clientIdentificationString,
                                     IdentificationString serverIdentificationString,
                                     SshMsgKexInit clientKexInitMessage,
                                     SshMsgKexInit serverKexInitMessage,
                                     byte[] hostKeyBlob,
                                     PublicKey clientKey,
                                     PublicKey serverKey,
                                     byte[] sharedSecret)
            throws Exception {
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        StringCodec.writeValue(clientIdentificationString.rawData, buffer);
        StringCodec.writeValue(serverIdentificationString.rawData, buffer);
        StringCodec.writeValue(clientKexInitMessage.rawData, buffer);
        StringCodec.writeValue(serverKexInitMessage.rawData, buffer);
        StringCodec.writeValue(hostKeyBlob, buffer);
        EcPublicKeyCodec.writeValue(clientKey, buffer);
        EcPublicKeyCodec.writeValue(serverKey, buffer);
        buffer.write(sharedSecret); // already in mpint format

        return hash.reset().update(buffer.buffer(), 0, buffer.position()).compute();
    }

    protected byte[] hostKeyBlobFor(PublicKey hostKey)
            throws Exception {
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        DataCodec<PublicKey> publicKeyCodec = this.provider instanceof String
                                              ? PublicKeyCodecs.init(new PublicKeyCodec(), (String) this.provider)
                                              : this.provider instanceof Provider ? PublicKeyCodecs.init(new PublicKeyCodec(), (Provider) this.provider)
                                                                                  : PublicKeyCodecs.init(new PublicKeyCodec(), (String) null);
        publicKeyCodec.write(hostKey, buffer);

        return buffer.toByteArray();
    }
}
