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.marshaling.primitive.Uint32Codec;
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 DiffieHellmanGroupKeyExchange {

    protected static final String KEY_ALGORITHM = "DH";

    protected static final int MAX_KEY_GENERATION_ATTEMPTS = 1000;

    protected final MessageCodec.Factory codecFactory;
    protected final String hashAlgorithm;
    protected final Object provider;

    public DiffieHellmanGroupKeyExchange(MessageCodec.Factory codecFactory, String hashAlgorithm, Object provider) {
        assertNotNull(codecFactory);
        this.codecFactory = codecFactory;

        assertNotNull(hashAlgorithm);
        this.hashAlgorithm = hashAlgorithm;

        this.provider = provider;
    }

    @Override
    public String toString() {
        return "DiffieHellmanGroupKeyExchange {"
                + "hashAlgorithm='" + this.hashAlgorithm + '\''
                + ", provider='" + this.provider + '\''
                + '}';
    }

    protected KeyPair generateKeyPair(BigInteger p, BigInteger g)
            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(p, g));

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

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

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

    protected boolean keyPairValid(KeyPair keyPair, BigInteger p) {
        BigInteger e = ((DHPublicKey) keyPair.getPublic()).getY();
        return e.compareTo(BigInteger.ONE) > 0 && e.compareTo(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, BigInteger p)
            throws Exception {
        BigInteger sharedSecret = remoteKey.modPow(localKey, 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,
                                     Integer min,
                                     Integer n,
                                     Integer max,
                                     BigInteger p,
                                     BigInteger g,
                                     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);

        if (min != null) {
            Uint32Codec.writeValue(min, buffer);
        }

        Uint32Codec.writeValue(n, buffer);

        if (max != null) {
            Uint32Codec.writeValue(max, buffer);
        }

        MpintCodec.writeValue(p, buffer);
        MpintCodec.writeValue(g, 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();
    }
}
