/*
 * Copyright (c) 2004-2017, by Alex Leigh.
 * All rights reserved.
 *
 * THIS IS UNPUBLISHED PROPRIETARY SOURCE CODE
 * The copyright notice above does not evidence any
 * actual or intended publication of such source code.
 */

package lc.mecha.exec;

import com.jcraft.jsch.*;
import lc.mecha.cred.Credential;
import lc.mecha.log.MechaLogger;
import lc.mecha.log.MechaLoggerFactory;
import lc.mecha.util.IOUtil;
import lc.mecha.util.StringAccumulatorV2;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.List;
import java.util.Properties;

/**
 * This class provides an executor service facade for SSH. The class is limited in that it will only function with
 * remote commands which do not require input.
 * <p>
 * Implementation notes:
 * <p>
 * For the duration this object is active, a SSH session will be maintained with the remote endpoint. If the session
 * fails, the class will not attempt to auto-connect and subsequent requests will fail.
 * <p>
 * Calling <code>exec()</code> multiple times gives reasonable performance since a new session does not have to be
 * created each time, only a channel. However, whatever state on the remote endpoint is maintained between
 * executions (for example, if a previous remote command has changed the remote execution environment or caused
 * an error).
 * <p>
 * Many remote endpoints have idle-timeouts. If this object is created and a session established, if an idle-timeout
 * occurs on the remote, the session will be in a disconnected state. As a result, further executions will fail,
 * as this class does not re-establish the connection automatically.
 *
 * @author Alex Leigh
 */
// TODO: Add retry mechanic
// TODO: Lazy initialize the JSSH channel
public class SshExecutor implements ProcessExecutor, AutoCloseable {
    private static final MechaLogger logger = MechaLoggerFactory.getLogger(SshExecutor.class);
    private JSch jsch = new JSch();
    private Session session;
    private final String host;
    private final int port;
    private final String username;
    private final String password;

    public SshExecutor(String host, int port, Credential cred) throws Exception {
        this(host, port, cred.getUsername(), cred.getPassword());
    }

    public SshExecutor(String host, int port, String username, String password) throws Exception {
        this.host = host;
        this.port = port;
        this.username = username;
        this.password = password;
        openSession();
    }

    /**
     * Open a session to the configured peer. If a session is already open, it will be closed and a new one
     * will be created.
     *
     * @throws JSchException
     */
    private void openSession() throws JSchException {
        if (session != null) {
            session.disconnect();
        }

        session = jsch.getSession(username, host, port);
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");
        session.setConfig(config);
        session.setPassword(password);
        session.connect();
    }

    /**
     * Transmit the given input stream to the SSH peer as a file.
     * <p>
     * Note: This method will not <code>close()</code> the <code>data</code> stream. If an exception is thrown,
     * some data may have been read from the stream.
     *
     * @param filename
     * @param data
     */
    public void upload(String filename, InputStream data, long length) throws Exception {
        // exec 'scp -t rfile' remotely

        boolean ptimestamp = false;

        String command = "scp " + (ptimestamp ? "-p" : "") + " -t " + filename;
        Channel channel = session.openChannel("exec");
        try {
            ((ChannelExec) channel).setCommand(command);

            // get I/O streams for remote scp
            OutputStream out = channel.getOutputStream();
            InputStream in = channel.getInputStream();
            channel.connect();

            if (checkAck(in) != 0) {
                throw new Exception("ACK failed");
            }

            long now = System.currentTimeMillis() / 1000;
            if (ptimestamp) {
                command = "T " + (now) + " 0";
                command += (" " + (now) + " 0\n");
                out.write(command.getBytes());
                out.flush();
                if (checkAck(in) != 0) {
                    System.exit(0);
                }
            }

            long filesize = length;
            command = "C0644 " + filesize + " ";
            if (filename.lastIndexOf('/') > 0) {
                command += filename.substring(filename.lastIndexOf('/') + 1);
            } else {
                command += filename;
            }
            command += "\n";
            out.write(command.getBytes());
            out.flush();
            if (checkAck(in) != 0) {
                System.exit(0);
            }

            IOUtil.copy(data, out);
            out.close();
            in.close();
        } finally {
            channel.disconnect();
        }
    }

    /**
     * Disconnect the SSH session and release any associated resources.
     */
    public void close() {
        session.disconnect();
    }

    private static int checkAck(InputStream in) throws IOException {
        int b = in.read();

        logger.debug("checkAck(): read {}", b);
        // b may be 0 for success,
        //          1 for error,
        //          2 for fatal error,
        //          -1
        if (b == 0) return b;
        if (b == -1) return b;

        if (b == 1 || b == 2) {
            StringBuffer sb = new StringBuffer();
            int c;
            do {
                c = in.read();
                sb.append((char) c);
            }
            while (c != '\n');
            if (b == 1) { // error
                System.out.print(sb.toString());
            }
            if (b == 2) { // fatal error
                System.out.print(sb.toString());
            }
        }
        return b;
    }

    @Override
    public ProcessExecution execute(String cmd, List<String> args) throws Exception {
        return execute(cmd, args.stream().toArray(String[]::new));
    }

    public ProcessExecution execute(String cmd) throws Exception {
        return execute(cmd, new String[]{});
    }

    @Override
    /**
     * Execute the given process on the remote machine. <code>args</code> may be <code>null</code>.
     *
     * If arguments are provided, each argument is surrounded by a double-quote character. Additional processing is
     * done to remove double quotes provided by the caller to prevent malicious arguments from ending the command
     * string prematurely with a dangling quote. This behavior can cause problems on Windows systems with
     * Powershell if one of the arguments needs to contain double quotes.
     *
     * No such processing is done on <code>cmd</code> however, so a full string with spaces and quotes may be
     * passed along with <code>null</code> arguments to result in the desired behavior. In this case, the caller
     * is responsible for determining if <code>cmd</code> is safe to execute on the remote system.
     */
    public ProcessExecution execute(String cmd, String[] args) throws Exception {
        if (!session.isConnected()) {
            throw new IllegalStateException("Session was found unexpectedly closed.");
        }

        ChannelExec channelExec = (ChannelExec) session.openChannel("exec");
        StringAccumulatorV2 c = new StringAccumulatorV2(" ");

        c.add(cmd);
        if (args != null) {
            for (String s : args) {
                // Prevent the caller from ending the string run early
                c.add("\"" + s.replace("\"", "\\\\") + "\"");
            }
        }

        String fulLCmd = c.asString();
        logger.debug("Executing command: {}", fulLCmd);
        channelExec.setCommand(fulLCmd);
        channelExec.connect();
        return new SshExecution(channelExec);
    }

    public long download(String filename, OutputStream data) throws Exception {
        long l = -1L;
        // exec 'scp -f rfile' remotely

        String command = "scp -f " + filename;
        Channel channel = session.openChannel("exec");
        try {
            ((ChannelExec) channel).setCommand(command);

            // get I/O streams for remote scp
            OutputStream out = channel.getOutputStream();
            InputStream in = channel.getInputStream();
            channel.connect();

            byte[] buf = new byte[1024];

            // send '\0'
            buf[0] = 0;
            out.write(buf, 0, 1);
            out.flush();

            while (true) {
                int c = checkAck(in);
                if (c != 'C') {
                    break;
                }

                // read '0644 '
                in.read(buf, 0, 5);

                long filesize = 0L;
                while (true) {
                    if (in.read(buf, 0, 1) < 0) {
                        // error
                        break;
                    }
                    if (buf[0] == ' ') break;
                    filesize = filesize * 10L + (long) (buf[0] - '0');
                }

                String file = null;
                for (int i = 0; ; i++) {
                    in.read(buf, i, 1);
                    if (buf[i] == (byte) 0x0a) {
                        file = new String(buf, 0, i);
                        break;
                    }
                }

                l = filesize;
                logger.info("filesize=" + filesize + ", file=" + file);

                // send '\0'
                buf[0] = 0;
                out.write(buf, 0, 1);
                out.flush();

                // read a content of lfile
                int foo;
                while (true) {
                    if (buf.length < filesize) foo = buf.length;
                    else foo = (int) filesize;
                    foo = in.read(buf, 0, foo);
                    if (foo < 0) {
                        // error
                        break;
                    }
                    data.write(buf, 0, foo);
                    filesize -= foo;
                    if (filesize == 0L) break;
                }

                if (checkAck(in) != 0) {
                    System.exit(0);
                }

                // send '\0'
                buf[0] = 0;
                out.write(buf, 0, 1);
                out.flush();
            }

            out.close();
            in.close();
        } finally {
            channel.disconnect();
        }

        return l;
    }
}