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.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.SshMsgKexEcdhInit;
import com.jscape.inet.ssh.protocol.v2.messages.SshMsgKexEcdhReply;
import com.jscape.util.net.connection.SyncMessageConnection;

import java.security.KeyPair;
import java.security.Provider;
import java.security.PublicKey;

/**
 * @author Alex Uzun
 */
public class EcDiffieHellmanClientKeyExchange
        extends EcDiffieHellmanKeyExchange
        implements ClientKeyExchange {

    protected final boolean hostKeySignatureVerificationRequired;

    public EcDiffieHellmanClientKeyExchange(String curveIdentifier, boolean hostKeySignatureVerificationRequired, Object provider) {
        super(curveIdentifier, 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();

        connection.write(new SshMsgKexEcdhInit(keyPair.getPublic()));
        SshMsgKexEcdhReply reply = (SshMsgKexEcdhReply) connection.read();

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

        byte[] sharedSecret = computeSharedSecret(keyPair, reply.qs);
        Hash hash = hashFor(keyPair.getPublic());
        byte[] exchangeHash = exchangeHashFor(hash,
                                              session.clientIdentificationString,
                                              session.serverIdentificationString,
                                              session.clientKexInitMessage,
                                              session.serverKexInitMessage,
                                              reply.hostKeyBlob,
                                              keyPair.getPublic(),
                                              reply.qs,
                                              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);
    }

}
