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.*;
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;

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

/**
 * @author Alex Usun
 */
public class DiffieHellmanGroupClientKeyExchange
        extends DiffieHellmanGroupKeyExchange
        implements ClientKeyExchange {

    public static final int DEFAULT_MIN = DiffieHellmanGroup.GROUP_1.p.bitLength();
    public static final int DEFAULT_N = DiffieHellmanGroup.GROUP_14.p.bitLength();
    public static final int DEFAULT_MAX = DiffieHellmanGroup.GROUP_18.p.bitLength();

    private static final int MIN_GROUP_SIZE = 1024;
    private static final int MAX_GROUP_SIZE = 8192;

    private final int min;
    private final int n;
    private final int max;
    private final boolean hostKeySignatureVerificationRequired;

    public static DiffieHellmanGroupClientKeyExchange defaultExchange(MessageCodec.Factory codecFactory, String hashAlgorithm, boolean hostKeySignatureVerificationRequired, String provider) {
        return new DiffieHellmanGroupClientKeyExchange(codecFactory, hashAlgorithm, provider, DEFAULT_MIN, DEFAULT_N, DEFAULT_MAX, hostKeySignatureVerificationRequired);
    }

    public static DiffieHellmanGroupClientKeyExchange defaultExchange(MessageCodec.Factory codecFactory, String hashAlgorithm, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return new DiffieHellmanGroupClientKeyExchange(codecFactory, hashAlgorithm, provider, DEFAULT_MIN, DEFAULT_N, DEFAULT_MAX, hostKeySignatureVerificationRequired);
    }

    public static DiffieHellmanGroupClientKeyExchange defaultExchange(String hashAlgorithm, boolean hostKeySignatureVerificationRequired, String provider) {
        return defaultExchange(DiffieHellmanKeyGroupExchangeMessages.defaultMessages(provider), hashAlgorithm, hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanGroupClientKeyExchange defaultExchange(String hashAlgorithm, boolean hostKeySignatureVerificationRequired, Provider provider) {
        return defaultExchange(DiffieHellmanKeyGroupExchangeMessages.defaultMessages(provider), hashAlgorithm, hostKeySignatureVerificationRequired, provider);
    }

    public static DiffieHellmanGroupClientKeyExchange defaultExchange(String hashAlgorithm, boolean hostKeySignatureVerificationRequired) {
        return defaultExchange(hashAlgorithm, hostKeySignatureVerificationRequired, CryptoServicesProvider.instance());
    }

    public DiffieHellmanGroupClientKeyExchange(MessageCodec.Factory codecFactory, String hashAlgorithm, Object provider, int min, int n, int max, boolean hostKeySignatureVerificationRequired) {
        super(codecFactory, hashAlgorithm, provider);

        assertTrue(MIN_GROUP_SIZE <= min && min <= n && n <= max && max <= MAX_GROUP_SIZE, "Bad group size values.");
        this.min = min;
        this.n = n;
        this.max = max;

        this.hostKeySignatureVerificationRequired = hostKeySignatureVerificationRequired;
    }

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

        connection.write(new SshMsgKexDhGexRequest(this.min, this.n, this.max));
        SshMsgKexDhGexGroup groupMessage = (SshMsgKexDhGexGroup) connection.read();

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

        connection.write(new SshMsgKexDhGexInit(e));
        SshMsgKexDhGexReply reply = (SshMsgKexDhGexReply) connection.read();

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

        byte[] sharedSecret = computeSharedSecret(x, reply.f, groupMessage.p);
        Hash hash = hash();
        byte[] exchangeHash = exchangeHashFor(hash,
                                              session.clientIdentificationString,
                                              session.serverIdentificationString,
                                              session.clientKexInitMessage,
                                              session.serverKexInitMessage,
                                              reply.hostKeyBlob,
                                              this.min,
                                              this.n,
                                              this.max,
                                              groupMessage.p,
                                              groupMessage.g,
                                              e,
                                              reply.f,
                                              sharedSecret);

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

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

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

    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);
    }

}
