package com.jscape.inet.ssh.protocol.v2.authentication;

import com.jscape.inet.ssh.protocol.messages.Message;
import com.jscape.inet.ssh.protocol.v2.marshaling.MessageCodec;
import com.jscape.inet.ssh.protocol.v2.messages.*;
import com.jscape.inet.ssh.protocol.v2.transport.TransportConnection;

import java.security.KeyPair;

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

/**
 * Represents a SSH 'public-key' authentication operation.
 *
 * @author Alex Usun
 */
public class PublicKeyClientAuthentication
        extends AbstractClientAuthentication {

    private static final SignatureService SIGNATURE_SERVICE = SignatureService.defaultService();

    private final String username;
    private final PublicKeyAlgorithm publicKeyAlgorithm;
    private final KeyPair keys;

    public static PublicKeyClientAuthentication defaultAuthentication(String username, PublicKeyAlgorithm publicKeyAlgorithm, KeyPair keys) {
        return new PublicKeyClientAuthentication(PublicKeyAuthenticationMessages.defaultMessages(), username, publicKeyAlgorithm, keys);
    }

    /**
     * Creates a new operation instance.
     *
     * @param username an username
     * @param keys key pair
     */
    public PublicKeyClientAuthentication(MessageCodec.Factory codecFactory, String username, PublicKeyAlgorithm publicKeyAlgorithm, KeyPair keys) {
        super(codecFactory);

        assertNotNull(username);
        this.username = username;

        assertNotNull(publicKeyAlgorithm);
        this.publicKeyAlgorithm = publicKeyAlgorithm;

        assertNotNull(keys);
        this.keys = keys;
    }

    @Override
    public void applyTo(TransportConnection connection, String serviceName)
            throws Exception {
        init(connection);
        sendPublicKey(connection, serviceName);
        sendSignature(connection, serviceName);
    }

    @Override
    public String toString() {
        return "PublicKeyClientAuthentication {"
               + "username='" + this.username + '\''
               + ", keys=" + this.keys
               + ", banner='" + this.banner + '\''
               + ", bannerLanguageTag='" + this.bannerLanguageTag + '\''
               + '}';
    }

    protected void init(TransportConnection connection) {
        super.init(connection);
        PublicKeyAuthenticationMessages.init(connection.session().messageCodec);
    }

    private void sendPublicKey(TransportConnection connection, String serviceName)
            throws Exception {
        connection.write(new SshMsgUserAuthRequestPublicKey(this.username, serviceName, this.publicKeyAlgorithm.identifier, this.keys.getPublic()));
        Message response = receive(connection);
        assertPublicKeyOk(response);
    }

    private void sendSignature(TransportConnection connection, String serviceName)
            throws Exception {
        Signature signature = SIGNATURE_SERVICE.signatureFor(connection.session().sessionId, this.username, serviceName, this.publicKeyAlgorithm.identifier, this.keys);
        connection.write(new SshMsgUserAuthRequestPublicKeySignature(this.username, serviceName, this.publicKeyAlgorithm.identifier, this.keys.getPublic(), signature));
        Message response = receive(connection);
        assertSuccess(response);
    }

    private void assertPublicKeyOk(Message response)
            throws InvalidCredentialsException {
        if (!(response instanceof SshMsgUserAuthPkOk)) {
            throw new InvalidCredentialsException();
        }
    }

}
