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.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.SshMsgKexInit;
import com.jscape.util.io.ByteArrayOutputStream;
import com.jscape.util.io.DataCodec;

import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.Provider;
import java.security.PublicKey;

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

/**
 * @author Alex Usun
 */
public abstract class DiffieHellmanKeyExchange {

    protected static final String KEY_ALGORITHM = "DH";

    protected static final int MAX_KEY_GENERATION_ATTEMPTS = 1000;

    protected final MessageCodec.Factory codecFactory;
    protected final BigInteger p;
    protected final BigInteger g;
    protected final String hashAlgorithm;
    protected final Object provider;

    public DiffieHellmanKeyExchange(MessageCodec.Factory codecFactory, BigInteger p, BigInteger g, String hashAlgorithm, Object provider) {
        assertNotNull(codecFactory);
        this.codecFactory = codecFactory;

        assertNotNull(p);
        this.p = p;

        assertNotNull(g);
        this.g = g;

        assertNotNull(hashAlgorithm);
        this.hashAlgorithm = hashAlgorithm;

        this.provider = provider;
    }

    @Override
    public String toString() {
        return "DiffieHellmanKeyExchange {"
               + "codecFactory=" + this.codecFactory
               + ", p=" + this.p
               + ", g=" + this.g
               + ", hashAlgorithm='" + this.hashAlgorithm + '\''
               + ", provider=" + this.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 DHParameterSpec(this.p, this.g));

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

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

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

    protected boolean keyPairValid(KeyPair keyPair) {
        BigInteger e = ((DHPublicKey) keyPair.getPublic()).getY();
        return e.compareTo(BigInteger.ONE) > 0 && e.compareTo(this.p.subtract(BigInteger.ONE)) < 0;
    }

    protected JceHash hash()
            throws Exception {
        return this.provider instanceof String
               ? JceHash.hashFor(this.hashAlgorithm, (String) this.provider)
               : this.provider instanceof Provider ? JceHash.hashFor(this.hashAlgorithm, (Provider) this.provider)
                                                   : JceHash.hashFor(this.hashAlgorithm, (String) null);
    }

    protected byte[] computeSharedSecret(BigInteger localKey, BigInteger remoteKey)
            throws Exception {
        BigInteger sharedSecret = remoteKey.modPow(localKey, this.p);
        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,
                                     BigInteger e,
                                     BigInteger f,
                                     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);
        MpintCodec.writeValue(e, buffer);
        MpintCodec.writeValue(f, 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();
    }
}
