package be.iminds.ilabt.jfed.connectivity_tester;

import be.iminds.ilabt.jfed.lowlevel.GeniUserProvider;
import be.iminds.ilabt.jfed.lowlevel.connection.JFedConnection;
import be.iminds.ilabt.jfed.lowlevel.ssh_key_info.SshKeyInfo;
import be.iminds.ilabt.jfed.lowlevel.ssh_key_info.SshKeyInfoFactory;
import be.iminds.ilabt.jfed.preferences.JFedGuiPreferences;
import be.iminds.ilabt.jfed.util.SshProxySocketFactory;
import net.schmizz.sshj.SSHClient;
import net.schmizz.sshj.connection.channel.direct.Session;
import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
import net.schmizz.sshj.userauth.keyprovider.KeyPairWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.SocketImplFactory;
import java.util.stream.Collectors;

/**
 * User: twalcari
 * Date: 3/5/14
 * Time: 4:09 PM
 */
public class SSHLoginTest implements ConnectivityTest {
    private static final Logger LOG = LoggerFactory.getLogger(SSHLoginTest.class);
    private final String hostname;
    private final int port;
    private final String username;
    private final SshKeyInfo forceUserSshKeyInfo;
    private final JFedConnection.SshProxyInfo proxyInfo;

    //FIXME: how can we inject this value? We are loading this class dynamically...
    private final GeniUserProvider geniUserProvider;
    private final JFedGuiPreferences jFedPreferences;

    /**
     * @param forceUserSshKeyInfo if not null, this SSH key info will be used. If null, it will be derived from the preferences.
     * @param proxyInfo set to proxyInfo when needed. null when no proxy. Hint: guiModel.getSfaConnectionPool().getProxyInfo();
     */
    public SSHLoginTest(String hostname,
                        int port,
                        String username,
                        @Nullable SshKeyInfo forceUserSshKeyInfo,
                        JFedConnection.SshProxyInfo proxyInfo,
                        GeniUserProvider geniUserProvider,
                        JFedGuiPreferences jFedPreferences) {
        this.hostname = hostname;
        this.port = port;
        this.username = username;
        this.forceUserSshKeyInfo = forceUserSshKeyInfo;
        this.proxyInfo = proxyInfo;
        this.geniUserProvider = geniUserProvider;
        this.jFedPreferences = jFedPreferences;
    }

    @Override
    public String getName() {
        return "SSH Login " + username + "@" + hostname + ":" + port + (proxyInfo == null ? ""
                : " with proxy " + proxyInfo);
    }

    @Override
    public boolean isLoginRequired() {
        return true;
    }

    @Override
    public ConnectivityTestResult call() {
        SshKeyInfo sshKeyInfo;
        if (forceUserSshKeyInfo != null) {
            sshKeyInfo = forceUserSshKeyInfo;
            if (!forceUserSshKeyInfo.hasUnencryptedPrivateKey()) {
                LOG.info("Private key in forceUserSshKeyInfo is encrypted. Login cannot be tested." +
                        " forceUserSshKeyInfo.getPrivateKey()=" + (sshKeyInfo.getPrivateKey() == null ? "null" : "nonnull") +
                        " forceUserSshKeyInfo.hasUnencryptedPrivateKey()=" + sshKeyInfo.hasUnencryptedPrivateKey());
                return new ConnectivityTestResult(Status.SKIPPED, "Private key in forceUserSshKeyInfo is encrypted. Login cannot be tested.");
            }
        } else {
            sshKeyInfo = jFedPreferences.getCustomSshKeyInfo();
        }
        if (sshKeyInfo != null && !sshKeyInfo.hasUnencryptedPrivateKey()) {
            LOG.info("Private key in preferences is encrypted. Login cannot be tested." +
                    " sshKeyInfo.getPrivateKey()=" + (sshKeyInfo.getPrivateKey() == null ? "null" : "nonnull") +
                    " sshKeyInfo.hasUnencryptedPrivateKey()=" + sshKeyInfo.hasUnencryptedPrivateKey());
            return new ConnectivityTestResult(Status.SKIPPED, "Private key in preferences is encrypted. Login cannot be tested.");
        }
        if (sshKeyInfo == null) {
            LOG.info("The user's SSH key will be used.");
            final be.iminds.ilabt.jfed.lowlevel.GeniUser geniUser = geniUserProvider.getLoggedInGeniUser();
            sshKeyInfo = SshKeyInfoFactory.createGeniUserSshKeyInfo(geniUser);
        }
        if (!sshKeyInfo.hasUnencryptedPrivateKey() || sshKeyInfo.getPrivateKey() == null) {
            LOG.info("Private key in preferences is encrypted. Login cannot be tested." +
                            " sshKeyInfo.hasUnencryptedPrivateKey()={}" +
                            " sshKeyInfo.getPrivateKey()= {}",
                    sshKeyInfo.hasUnencryptedPrivateKey(),
                    sshKeyInfo.getPrivateKey() == null ? "null" : "notnull");
            return new ConnectivityTestResult(Status.SKIPPED, "Private key in preferences is encrypted (" + (
                    sshKeyInfo.hasUnencryptedPrivateKey() ? "false, but privKey=" + (
                            sshKeyInfo.getPrivateKey() == null ? "null" : "notnull") : "true"
            ) + "). Login cannot be tested.");
        }

        //Create SSH Proxy connection if necessary
        SSHClient ssh = new SSHClient();
        ssh.addHostKeyVerifier(new PromiscuousVerifier());
        SocketImplFactory socketImplFactory = null;
        boolean forceSshProxyDns = false;
        try {
            if (proxyInfo != null) {
                ssh.setSocketFactory(SshProxySocketFactory.createWithForcedTarget(proxyInfo, hostname, port));
                forceSshProxyDns = true;
                LOG.info("Using SSH proxy for SSH connection: " + proxyInfo);
            }
        } catch (IOException ex) {
            return new ConnectivityTestResult(Status.FAILED, "Could not connect to SSH proxy: " + ex.getMessage(), ex);
        }

        try {
            if (forceSshProxyDns)
                ssh.connect("localhost", port);
            else
                ssh.connect(hostname, port);
        } catch (IOException ex) {
            return new ConnectivityTestResult(Status.FAILED, "Could not connect to " + hostname + ": " + ex.getMessage(), ex);
        }

        try {

            ssh.authPublickey(username, new KeyPairWrapper(sshKeyInfo.getPublicKey(), sshKeyInfo.getPrivateKey()));
        } catch (IOException e) {
            LOG.error("Error while registring private key", e);
            return new ConnectivityTestResult(Status.FAILED, "Couldn't authenticate with provided private key", e);
        }

        LOG.debug("Successfully authenticated {}@{}:{}", username, hostname, port);


        try (Session session = ssh.startSession()) {
            Session.Command comm = session.exec("uname -a");
            comm.join();

            LOG.debug("Test-channel exited with status {}", comm.getExitStatus());

            BufferedReader sout = new BufferedReader(new InputStreamReader(comm.getInputStream()));
            BufferedReader serr = new BufferedReader(new InputStreamReader(comm.getErrorStream()));


            LOG.info("Output of test-commands: {}", sout.lines().collect(Collectors.joining("\n")));
            String errorOutput = serr.lines().collect(Collectors.joining("\n"));
            if (!errorOutput.isEmpty())
                LOG.warn("Error-output of test-commands: " + errorOutput);

            return new ConnectivityTestResult(Status.SUCCEEDED,
                    "Successfully authenticated " + username + "@" + hostname + ":" + port);
        } catch (IOException e) {
            return new ConnectivityTestResult(Status.FAILED,
                    "Successfully authenticated " + username + "@" + hostname + ":" + port + ". But command failed", e);
        } finally {
            try {
                ssh.close();
            } catch (IOException ignored) {
            }
        }
    }
}
