package be.iminds.ilabt.jfed.ssh_terminal_tool.unix;

import be.iminds.ilabt.jfed.util.IOUtils;
import be.iminds.ilabt.jfed.util.KeyUtil;
import org.slf4j.LoggerFactory;

import javax.inject.Inject;
import javax.inject.Singleton;
import java.io.*;
import java.security.PublicKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Arrays;
import java.util.Objects;

/**
 * OpenSshKnownHostsFile
 */
@Singleton
public class OpenSshKnownHostsFile {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(OpenSshKnownHostsFile.class);

    private final File knownHostsFile;

    /**
     * for testing
     */
    public OpenSshKnownHostsFile(File knownHostsFile) {
        this.knownHostsFile = knownHostsFile;
    }

    @Inject
    public OpenSshKnownHostsFile() {
        File userHomeDir = new File(System.getProperty("user.home"));
        File userSshDir = new File(userHomeDir, ".ssh");
        this.knownHostsFile = new File(userSshDir, "known_hosts");
    }

    /**
     * only adds when not yet added.
     * <p/>
     * On error, will no throw exception, but do nothing instead.
     */
    public void addKnownHost(String hostname, PublicKey publicKey) {
        if (!(publicKey instanceof RSAPublicKey)) {
            LOG.warn("publicKey must be RSAPublicKey: aborting adding host key");
            return;
        }

        RSAPublicKey rsaPubKey = (RSAPublicKey) publicKey;
        String knownHostFormatPublicKey = KeyUtil.rsaPublicKeyToOpenSshAuthorizedKeysFormatData(rsaPubKey);

        addKnownHost(hostname, knownHostFormatPublicKey);
    }

    /**
     * only adds when not yet added.
     * <p/>
     * On error, will no throw exception, but do nothing instead.
     */
    public void addKnownHost(String hostname, String origKnownHostFormatPublicKey) {
        //first handle extra spaces before and after, and extra spaces between parts
        String knownHostFormatPublicKey = origKnownHostFormatPublicKey.trim().replaceAll("[ \n\t]+", " ");

        //second check format of knownHostFormatPublicKey.
        //might be plain key, or might start with <hostname> ssh-rsa
        String[] knownHostFormatPublicKeyParts = knownHostFormatPublicKey.split(" ");
        //noinspection StatementWithEmptyBody
        if (knownHostFormatPublicKeyParts.length == 1) {
            //nothing to do
        } else if (knownHostFormatPublicKeyParts.length == 3 || knownHostFormatPublicKeyParts.length == 4) {
            knownHostFormatPublicKey = knownHostFormatPublicKeyParts[2];
            if (!Objects.equals(knownHostFormatPublicKeyParts[1], "ssh-rsa"))
                throw new RuntimeException("Unknown format of known host key (" + knownHostFormatPublicKeyParts.length + " parts, " +
                        "part 2 should be ssh-rsa but is " + knownHostFormatPublicKeyParts[1] + "): \"" + knownHostFormatPublicKey + "\"");
        } else {
            throw new RuntimeException("Unknown format of known host key (" + knownHostFormatPublicKeyParts.length + " parts): \"" + knownHostFormatPublicKey + "\"");
        }


        LOG.debug("addKnownHost called for " + hostname + " -> " + knownHostFormatPublicKey);
        if (!knownHostsFile.exists()) {
            LOG.warn("No ~/.ssh/known_hosts file found: aborting adding host key");
            return;
        }
        try {
            if (knowsHost(hostname, knownHostFormatPublicKey)) {
                LOG.warn("Host key is already added: aborting adding host key");
                return;
            }
            String knownHostsLine = hostname + " ssh-rsa " + knownHostFormatPublicKey;

            LOG.debug("addKnownHost is adding a known hosts line: " + knownHostsLine);

            String knownHostsContent = IOUtils.fileToString(knownHostsFile);
            boolean endsWithNewline = knownHostsContent.endsWith("\n");

            PrintWriter pw = new PrintWriter(new FileWriter(knownHostsFile, true/*append*/));
            pw.println((endsWithNewline ? "" : "\n") + knownHostsLine);
            pw.close();
        } catch (IOException e) {
            LOG.warn("Failed to add to known_hosts: " + e.getMessage(), e);
        }
    }

    public boolean knowsHost(String hostname, PublicKey publicKey) throws IOException {
        if (!(publicKey instanceof RSAPublicKey))
            throw new RuntimeException("Only RSAPublicKey supported for now");

        RSAPublicKey rsaPubKey = (RSAPublicKey) publicKey;
        String knownHostFormatPublicKey = KeyUtil.rsaPublicKeyToOpenSshAuthorizedKeysFormatData(rsaPubKey);

        return knowsHost(hostname, knownHostFormatPublicKey);
    }

    public boolean knowsHost(String hostname, String knownHostFormatPublicKey) throws IOException {
        String[] knownHostFormatPublicKeyParts = knownHostFormatPublicKey.split(" ");
        //noinspection StatementWithEmptyBody
        if (knownHostFormatPublicKeyParts.length == 1) {
            //nothing to do
        } else if (knownHostFormatPublicKeyParts.length == 3 || knownHostFormatPublicKeyParts.length == 4) {
            knownHostFormatPublicKey = knownHostFormatPublicKeyParts[2];
            if (!Objects.equals(knownHostFormatPublicKeyParts[1], "ssh-rsa"))
                throw new RuntimeException("Unknown format of known host key (" + knownHostFormatPublicKeyParts.length + " parts, " +
                        "part 2 should be ssh-rsa but is " + knownHostFormatPublicKeyParts[1] + "): \"" + knownHostFormatPublicKey + "\"");
        } else {
            throw new RuntimeException("Unknown format of known host key (" + knownHostFormatPublicKeyParts.length + " parts): \"" + knownHostFormatPublicKey + "\"");
        }


        //main method: ask ssh-keygen
        try {
            String sshKeyGenRes = myCommandHelper(true, "ssh-keygen", "-F", hostname, "-f", knownHostsFile.getPath());
            if (sshKeyGenRes != null) {
                //output: empty when no keys found
                //        lines from known_hosts otherwise
                return sshKeyGenRes.contains(knownHostFormatPublicKey);
            }
        } catch (Exception e) {
            LOG.warn("Error while calling ssh-keygen: " + e.getMessage(), e);
            //ignore and go to fallback
        }

        LOG.warn("Using flawed callback to determine if host is in known_hosts");
        //fallback: search in content itself. This is flawed, as the host key might be registered only for other aliases of the host (for example: internal IP)
        BufferedReader br = null;
        try {
            br = new BufferedReader(new FileReader(knownHostsFile));
            String line = br.readLine();
            while (line != null) {
                if (line.contains(knownHostFormatPublicKey) && line.contains(hostname))
                    return true;
                line = br.readLine();
            }
        } finally {
            if (br != null)
                br.close();
        }
        return false;
    }


    private static String myCommandHelper(boolean mayFail, String... command) {

        LOG.debug("Starting command: " + Arrays.asList(command));

        ProcessBuilder pb = new ProcessBuilder(command);
        pb.redirectErrorStream();

        try {

            Process p = pb.start();
            p.getOutputStream().close(); //we don't write to the command's stdin
            BufferedReader input = new BufferedReader(new InputStreamReader(p.getInputStream()));

            StringBuilder res = new StringBuilder();
            String line;
            while ((line = input.readLine()) != null) {
                res.append(line);
            }
            //note: input might close before terminal stops.
            input.close();

            try {
                p.waitFor();
            } catch (InterruptedException e) {
                LOG.warn("SshAgentHelper.myCommandHelper thread interrupted while waiting for terminal to stop", e);
            }

            if (p.exitValue() != 0) {
                if (mayFail) {
                    LOG.info("SshAgentHelper command \"" + command[0] + "\" returned non zero exit value: " + p.exitValue() + " output: " + res.toString());
                    return null;
                } else
                    throw new RuntimeException("SshAgentHelper command \"" + command[0] + "\" returned non zero exit value: " + p.exitValue() + " output: " + res.toString());
            }

            return res.toString();
        } catch (IOException e) {
            LOG.error("Error executing command: " + e.getMessage(), e);
            return null;
        }

    }
}
