package utils;

import java.io.ByteArrayInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;

import net.sf.expectit.Expect;
import net.sf.expectit.ExpectBuilder;
import net.sf.expectit.matcher.Matchers;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.ChannelSftp;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.SftpException;

/**
* These 2 dependencies are required to run/compile this class:
*	<dependency>
*		<groupId>com.jcraft</groupId>
*		<artifactId>jsch</artifactId>
*		<version>0.1.51</version>
*	</dependency>
*	<dependency>
*		<groupId>net.sf.expectit</groupId>
*		<artifactId>expectit-core</artifactId>
*		<version>0.5.0</version>
*	</dependency>
**/
public class SSH {

    public enum ChannelType {

        SHELL, SFTP;

        public boolean isSftp() {
            return this.equals(SFTP);
        }
    };

    private static final String HOST = "HOST";
    private static final String USERNAME = "USERNAME";
    private static final String PASSWORD = "PASSWD";
    public static final String FILENAME = "FILE";

    private Session session;
    private Channel channel;
    private ChannelType channelType = ChannelType.SHELL;
    private Expect expect;

    public SSH() throws JSchException, IOException {
        initialize();
    }

    public SSH(ChannelType channelType) throws JSchException, IOException {
        this.channelType = channelType;
        initialize();
    }

    public String runCmd(String cmd, boolean... fetchResult) throws JSchException, IOException {
        if (channelType.isSftp()) {
            throw new UnsupportedOperationException("For the RUN operation, ChannelType must be SHELL.");
        }

        reconnectSSHIfNecessary();
        expect.sendLine(cmd).expect(Matchers.contains("$"));
        return fetchResult.length > 0 && fetchResult[0] ? printResult() : "";
    }

    public boolean put(byte[] file) throws FileNotFoundException, SftpException {
        if (!channelType.isSftp()) {
            throw new UnsupportedOperationException("For the PUT operation, ChannelType must be SFTP.");
        }

        ByteArrayInputStream bis = new ByteArrayInputStream(file);
        ((ChannelSftp) channel).put(bis, FILENAME + ".tar.gz");

        return true;
    }

    public void disconnect() throws JSchException, IOException {
        if (channelType.isSftp()) {
            ((ChannelSftp) channel).exit();
        } else {
            expect.sendLine("exit");
            expect.close();
        }

        channel.disconnect();
        session.disconnect();
    }

    private void initialize() throws JSchException, IOException {
        JSch jSch = new JSch();
        session = jSch.getSession(USERNAME, HOST);
        session.setPassword(PASSWORD);
        session.setConfig("StrictHostKeyChecking", "no");
        session.connect();

        channel = session.openChannel(channelType.toString().toLowerCase());
        expect = new ExpectBuilder()
                .withOutput(channel.getOutputStream())
                .withInputs(channel.getInputStream(), channel.getExtInputStream())
                .build();

        channel.connect();

        if (!channelType.isSftp()) {
            runCmd("stty -echo");
        }
    }

    private void reconnectSSHIfNecessary() throws JSchException, IOException {
        if (!session.isConnected() || channel.isClosed()) {
            initialize();
        }
    }

    private String printResult() throws IOException {
        expect.sendLine("exit");
        InputStream in = channel.getInputStream();
        int c;
        String result = "";
        while ((c = in.read()) != -1) {
            result += (char) c;
        }

        return result.substring(0, result.length());
    }

}