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

import com.jscape.crypto.Hash;
import com.jscape.inet.ssh.protocol.messages.Message;
import com.jscape.inet.ssh.protocol.v2.marshaling.ClientSession;
import com.jscape.inet.ssh.protocol.v2.marshaling.MessageCodec;
import com.jscape.inet.ssh.protocol.v2.marshaling.keys.ComponentSigner;
import com.jscape.inet.ssh.protocol.v2.marshaling.keys.Signer;
import com.jscape.inet.ssh.protocol.v2.marshaling.keys.Signers;
import com.jscape.inet.ssh.protocol.v2.messages.Signature;
import com.jscape.inet.ssh.protocol.v2.messages.SshMsgKexDhInit;
import com.jscape.inet.ssh.protocol.v2.messages.SshMsgKexDhReply;
import com.jscape.util.crypto.CryptoServicesProvider;
import com.jscape.util.net.connection.SyncMessageConnection;

import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.Provider;
import java.security.PublicKey;

/**
 * @author Alex Usun
 */
public class DiffieHellmanClientKeyExchange
        extends DiffieHellmanKeyExchange
        implements ClientKeyExchange {

    protected final boolean hostKeySignatureVerificationRequired;

    public static DiffieHellmanClientKeyExchange group1Sha1(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_1.p, DiffieHellmanGroup.GROUP_1.g, "SHA1", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group1Sha1(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_1.p, DiffieHellmanGroup.GROUP_1.g, "SHA1", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group1Sha1(boolean hostKeySignatureVerificationRequired, String provider) {
        return group1Sha1(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group1Sha1(boolean hostKeySignatureVerificationRequired, Provider provider) {
        return group1Sha1(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group1Sha1(boolean hostKeySignatureVerificationRequired) {
        return group1Sha1(hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public static DiffieHellmanClientKeyExchange group14Sha1(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_14.p, DiffieHellmanGroup.GROUP_14.g, "SHA1", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha1(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_14.p, DiffieHellmanGroup.GROUP_14.g, "SHA1", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha1(boolean hostKeySignatureVerificationRequired, String provider) {
        return group14Sha1(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha1(boolean hostKeySignatureVerificationRequired, Provider provider) {
        return group14Sha1(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha1(boolean hostKeySignatureVerificationRequired) {
        return group14Sha1(hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public static DiffieHellmanClientKeyExchange group14Sha256(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_14.p, DiffieHellmanGroup.GROUP_14.g, "SHA256", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha256(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_14.p, DiffieHellmanGroup.GROUP_14.g, "SHA256", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha256(boolean hostKeySignatureVerificationRequired, String provider) {
        return group14Sha256(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha256(boolean hostKeySignatureVerificationRequired, Provider provider) {
        return group14Sha256(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group14Sha256(boolean hostKeySignatureVerificationRequired) {
        return group14Sha256(hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public static DiffieHellmanClientKeyExchange group15Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_15.p, DiffieHellmanGroup.GROUP_15.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group15Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_15.p, DiffieHellmanGroup.GROUP_15.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group15Sha512(boolean hostKeySignatureVerificationRequired, String provider) {
        return group15Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group15Sha512(boolean hostKeySignatureVerificationRequired, Provider provider) {
        return group15Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group15Sha512(boolean hostKeySignatureVerificationRequired) {
        return group15Sha512(hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public static DiffieHellmanClientKeyExchange group16Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_16.p, DiffieHellmanGroup.GROUP_16.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group16Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_16.p, DiffieHellmanGroup.GROUP_16.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group16Sha512(boolean hostKeySignatureVerificationRequired, String provider) {
        return group16Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group16Sha512(boolean hostKeySignatureVerificationRequired, Provider provider) {
        return group16Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group16Sha512(boolean hostKeySignatureVerificationRequired) {
        return group16Sha512(hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public static DiffieHellmanClientKeyExchange group17Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_17.p, DiffieHellmanGroup.GROUP_17.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group17Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_17.p, DiffieHellmanGroup.GROUP_17.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group17Sha512(boolean hostKeySignatureVerificationRequired, String provider) {
        return group17Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group17Sha512(boolean hostKeySignatureVerificationRequired, Provider provider) {
        return group17Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group17Sha512(boolean hostKeySignatureVerificationRequired) {
        return group17Sha512(hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public static DiffieHellmanClientKeyExchange group18Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_18.p, DiffieHellmanGroup.GROUP_18.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group18Sha512(MessageCodec.Factory codecFactory, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanClientKeyExchange(codecFactory, DiffieHellmanGroup.GROUP_18.p, DiffieHellmanGroup.GROUP_18.g, "SHA512", hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group18Sha512(boolean hostKeySignatureVerificationRequired, String provider) {
        return group18Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group18Sha512(boolean hostKeySignatureVerificationRequired, Provider provider) {
        return group18Sha512(DiffieHellmanKeyExchangeMessages.defaultMessages(provider), hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanClientKeyExchange group18Sha512(boolean hostKeySignatureVerificationRequired) {
        return group18Sha512(hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public DiffieHellmanClientKeyExchange(MessageCodec.Factory codecFactory, BigInteger p, BigInteger g, String hashAlgorithm, boolean hostKeySignatureVerificationRequired, Object provider) {
        super(codecFactory, p, g, hashAlgorithm, provider);

        this.hostKeySignatureVerificationRequired = hostKeySignatureVerificationRequired;
    }

    @Override
    public KeyExchangeResult exchangeKeys(SyncMessageConnection<Message> connection, ClientSession session, HostKeyVerificationService hostKeyVerificationService)
            throws Exception {
        this.codecFactory.update(session.messageCodec);

        KeyPair keyPair = generateKeyPair();
        BigInteger x = ((DHPrivateKey) keyPair.getPrivate()).getX();
        BigInteger e = ((DHPublicKey) keyPair.getPublic()).getY();

        connection.write(new SshMsgKexDhInit(e));
        SshMsgKexDhReply reply = (SshMsgKexDhReply) connection.read();

        hostKeyVerificationService.assertKeyValid(connection.remoteAddress().inetAddress(), reply.hostKey);

        byte[] sharedSecret = computeSharedSecret(x, reply.f);
        Hash hash = hash();
        byte[] exchangeHash = exchangeHashFor(hash,
                                              session.clientIdentificationString,
                                              session.serverIdentificationString,
                                              session.clientKexInitMessage,
                                              session.serverKexInitMessage,
                                              reply.hostKeyBlob,
                                              e,
                                              reply.f,
                                              sharedSecret);

        if (this.hostKeySignatureVerificationRequired) {
            verifySignature(exchangeHash, reply.hostKey, reply.signature);
        }

        return new KeyExchangeResult(hash, sharedSecret, exchangeHash, reply.hostKey);
    }

    protected void verifySignature(byte[] exchangeHash, PublicKey hostKey, Signature signature)
            throws Exception {
        Signer signer = this.provider instanceof String ? Signers.init(new ComponentSigner(), (String) this.provider) : this.provider instanceof Provider ? Signers.init(new ComponentSigner(), (Provider) this.provider) : Signers.init(new ComponentSigner(), (String) null);
        signer.assertSignatureValid(exchangeHash, hostKey, signature);
    }

}
