/*
 * @(#)SshSession.java
 *
 * Copyright (c) 2003 JSCAPE
 * 1147 S. 53rd Pl., Mesa, Arizona, 85206, U.S.A.
 * All rights reserved.
 *
 * This software is the confidential and proprietary information of
 * JSCAPE. ("Confidential Information").  You shall not disclose such
 * Confidential Information and shall use it only in accordance with
 * the terms of the license agreement you entered into with JSCAPE.
 */
package com.jscape.inet.ssh;

import com.jscape.inet.ssh.util.SshParameters;
import com.jscape.util.Time;

import java.io.IOException;
import java.io.OutputStream;
import java.io.PrintStream;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;

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

/**
 * Convenience class for automating communications with a SSH server.
 * <p/>
 * Example:<P>
 * <pre>
 * String shellPrompt = "$";
 *
 * // creates new SshShession instance providing hostname, username and password arguments
 * SshSession session = new SshSession(hostname,username,password);
 *
 * // set expected shell prompt from SSH server
 * session.setShellPrompt(shellPrompt);
 *
 * // establish connection
 * session.connect();
 *
 * // send command to SSH shell and wait for shell prompt
 * session.send("cd /user/logs");
 *
 * // send command to SSH shell and wait for shell prompt
 * session.send("rm *.log");
 *
 * // send command to SSH shell and DO NOT wait for shell prompt
 * session.sendNoWait("exit");
 *
 * // close connection with SSH server
 * session.disconnect();
 * </pre>
 */
public class SshSession {

    private static final String DEFAULT_SHELL_PROMPT = "$";
    private static final String DEFAULT_EOL = "\n";
    private static final int DEFAULT_POLLING_INTERVAL = 100;
    private static final int DEFAULT_REGEX_MATCH_LENGTH = 20;

    private Ssh ssh;
    private String promptPattern;
    private boolean promptPatternIsRegex;
    private int regexMatchLength;
    private String endOfLine;
    private String characterSet;
    private int pollingInterval;
    private volatile boolean interrupted;
    private OutputStream out;

    public SshSession(Ssh ssh, String promptPattern, boolean promptPatternIsRegex, int regexMatchLength, String endOfLine, String characterSet, int pollingInterval) {
        assertNotNull(ssh);
        this.ssh = ssh;

        assertNotNull(promptPattern);
        this.promptPattern = promptPattern;

        this.promptPatternIsRegex = promptPatternIsRegex;
        this.regexMatchLength = regexMatchLength;

        assertNotNull(endOfLine);
        this.endOfLine = endOfLine;

        this.characterSet = characterSet;

        this.pollingInterval = pollingInterval;
    }

    /**
     * Creates a new SshSession instance.
     *
     * @param sshParams - the SshParameters instance.
     * @param ssh       - the Ssh instance.
     *
     * @throws SshException - the exception instance that is thrown if an error occurs.
     */
    public SshSession(SshParameters sshParams, Ssh ssh) {
        this(ssh, DEFAULT_SHELL_PROMPT, false, DEFAULT_REGEX_MATCH_LENGTH, DEFAULT_EOL, null, DEFAULT_POLLING_INTERVAL);
    }

    /**
     * Creates a new SshSession instance.
     *
     * @param sshParams an SshParameters instance
     *
     * @throws SshException if an SSH related exception occurs
     */
    public SshSession(SshParameters sshParams) {
        this(sshParams, new Ssh(sshParams));
    }

    /**
     * Creates a new SshSession instance.
     *
     * @param hostname the hostname of SSH server
     * @param port     the port of SSH server
     * @param username the SSH username
     * @param password the SSH password
     *
     * @throws SshException if an SSH related exception occurs
     */
    public SshSession(String hostname, int port, String username, String password) {
        this(new SshParameters(hostname, port, username, password));
    }

    /**
     * Creates a new SshSession instance.
     *
     * @param hostname the hostname of SSH server
     * @param username the SSH username
     * @param password the SSH password
     *
     * @throws SshException if an SSH related exception occurs
     */
    public SshSession(String hostname, String username, String password) {
        this(hostname, 22, username, password);
    }

    /**
     * Gets underlying <code>Ssh</code> instance.
     *
     * @return underlying <code>Ssh</code> instance
     */
    public synchronized Ssh getSsh() {
        return this.ssh;
    }

    /**
     * Sets the underlying <code>Ssh</code> instance.
     *
     * @param value
     */
    public synchronized void setSsh(Ssh value) {
        assertNotNull(value);
        this.ssh = value;
    }

    /**
     * Gets shell prompt to wait for before sending commands.
     * Default value is "$"
     *
     * @return the shell prompt
     */
    public synchronized String getShellPrompt() {
        return this.promptPattern;
    }

    /**
     * Sets shell prompt to wait for before sending commands.
     * Default value is "$"
     *
     * @param string the shell prompt
     */
    public synchronized void setShellPrompt(String string) {
        setShellPrompt(string, false);
    }

    /**
     * Sets shell prompt to wait for before sending commands.
     * Default value is "$"
     *
     * @param value the shell prompt
     * @param isRegex  true if shell prompt to match is a regular expression, false otherwise
     *
     * @see #setRegexMatchLength
     */
    public synchronized void setShellPrompt(String value, boolean isRegex) {
        assertNotNull(value);
        this.promptPattern = value;

        this.promptPatternIsRegex = isRegex;
    }

    /**
     * Gets terminator used when sending data to SSH server.
     * Default value is LF (Line Feed).
     *
     * @return the command terminator
     */
    public synchronized String getCommandTerminator() {
        return this.endOfLine;
    }

    /**
     * Sets terminator used when sending data to SSH server.
     * Default value is LF (Line Feed).
     *
     * @param value the command terminator
     */
    public synchronized void setCommandTerminator(String value) {
        assertNotNull(value);
        this.endOfLine = value;
    }

    /**
     * Returns echo flag.
     *
     * @return <code>true</code> if server echo is enabled;
     *         <code>false</code> otherwise
     */
    public synchronized boolean getEcho() {
        return this.ssh.getEcho();
    }

    /**
     * Sets server echo flag for the next connection.
     *
     * @param value if <code>true</code> server echo will
     *             be enabled for the next connection.
     */
    public synchronized void setEcho(boolean value) {
        this.ssh.setEcho(value);
    }

    /**
     * Gets the number of characters to read from the end of the response buffer when using regular expressions
     * to detect start or end prompts.  Default value is 20.
     *
     * @return match length
     */
    public synchronized int getRegexMatchLength() {
        return this.regexMatchLength;
    }

    /**
     * Sets the number of characters to read from the end of the response buffer when using regular expressions
     * to detect start or end prompts.  Default value is 20.
     *
     * @param value
     */
    public synchronized void setRegexMatchLength(int value) {
        this.regexMatchLength = value;
    }

    /**
     * Gets the character set
     *
     * @return The character set used by the service
     */
    public synchronized String getCharacterSet() {
        return this.characterSet;
    }

    /**
     * Sets the session character set
     *
     * @param value the character set
     */
    public synchronized void setCharacterSet(String value) {
        this.characterSet = value;
    }

    /**
     * Gets the polling interval used when checking for completion of commands sent using send or sendWait methods.
     * Default interval is 100ms.
     *
     * @return the polling interval in milliseconds.
     *
     * @see #send
     * @see #sendWait
     */
    public synchronized int getPollingInterval() {
        return this.pollingInterval;
    }

    /**
     * Sets the polling interval used when checking for completion of commands sent using send or sendWait methods.
     * Default interval is 100ms.
     *
     * @param value the polling interval in milliseconds.
     *
     * @see #send
     * @see #sendWait
     */
    public synchronized void setPollingInterval(int value) {
        this.pollingInterval = value;
    }

    /**
     * Checks if debugging is enabled.
     * Default is disabled.
     *
     * @return true if debugging enabled, false otherwise
     */
    public synchronized boolean getDebug() {
        return this.ssh.getDebug();
    }

    /**
     * Enables or disables debugging information.
     * In debug enabled mode all data received from SSH server is output to
     * the debug stream.  Default is disabled.
     *
     * @param value true to enable debugging, false to disable
     *
     * @see #setDebugStream
     */
    public synchronized void setDebug(boolean value) {
        this.ssh.setDebug(value);
    }

    /**
     * Sets PrintStream used in outputting debug information.
     * Default is System.out
     *
     * @param stream a PrintStream
     */
    public synchronized void setDebugStream(PrintStream stream) {
        this.ssh.setDebugStream(stream);
    }

    /**
     * Gets PrintStream used in outputting debug information.
     * Default is Sytem.out
     *
     * @return a PrintStream
     *
     * @see #setDebug
     * @see #setDebugStream
     */
    public synchronized PrintStream getDebugStream() {
        return this.ssh.getDebugStream();
    }

    /**
     * Add Ssh event listener.
     *
     * @param listener a SshListener
     *
     * @see SshListener
     */
    public synchronized void addSshListener(SshListener listener) {
        this.ssh.addSshListener(listener);
    }

    /**
     * Remove Ssh event listener.
     *
     * @param listener a SshListener
     *
     * @see SshListener
     */
    public synchronized void removeSshListener(SshListener listener) {
        this.ssh.removeSshListener(listener);
    }

    /**
     * Establishes connection with SSH server and performs login
     *
     * @param timeout timeout in milliseconds for establishing a connection
     *
     * @throws SshException if an SSH related exception occurs
     */
    public synchronized void connect(int timeout)
            throws SshException {
        SshTask task = taskFor(this.promptPattern, null, this.promptPattern, this.promptPatternIsRegex, timeout);
        SshScript script = new SshScript(this.ssh);
        script.addTask(task);

        if (timeout > 0) {
            this.ssh.setTimeout(timeout);
        }

        this.ssh.connect();
        this.out = this.ssh.getOutputStream();

        try {
            awaitTaskCompleted(task);
        } finally {
            cleanup(script);
        }
    }

    /**
     * Establishes connection with SSH server and performs login.
     *
     * @throws SshException if an SSH related exception occurs
     */
    public synchronized void connect()
            throws SshException {
        connect(0);
    }

    /**
     * Disconnects from SSH server.
     */
    public synchronized void disconnect() {
        this.ssh.disconnect();
    }

    /**
     * Sends data to SSH host without command terminator.
     * Does not wait for shell prompt to be returned.
     *
     * @param data the data to send
     *
     * @throws IOException if an I/O related error occurs
     * @see #setCommandTerminator
     */
    public synchronized void sendRaw(String data)
            throws IOException {
        this.out.write(this.characterSet != null ? data.getBytes(Charset.forName(this.characterSet)) : data.getBytes(StandardCharsets.UTF_8));
        this.out.flush();
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Does not wait for shell prompt to be returned.
     *
     * @param command the command to send
     *
     * @throws IOException if an I/O related error occurs
     * @see #setCommandTerminator
     */
    public synchronized void sendNoWait(String command)
            throws IOException {
        sendRaw(command + this.endOfLine);
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Waits until specified prompt is returned by SSH server buffering response data. If expected prompt
     * is not received within defined timeout then an SshTaskTimeoutException will be thrown.
     *
     * @param command the command to send
     * @param prompt  the prompt to wait for
     * @param timeout the timeout to use when waiting for task end prompt
     * @param regex   true if the prompt to match is a regular expression, false otherwise
     *
     * @return the response from telnet server
     *
     * @throws SshException if regular expression prompts are used and regular expression is invalid
     * @see #setCommandTerminator
     * @see #setRegexMatchLength
     * @see SshTaskTimeoutException
     * @see SshException
     */
    public synchronized String sendWait(String command, String prompt, boolean regex, long timeout)
            throws SshException {
        SshTask task = taskFor(null, command, prompt, regex, timeout);
        SshScript script = new SshScript(this.ssh);
        script.addTask(task);

        try {
            awaitTaskCompleted(task);
            return task.getResponse();
        } finally {
            cleanup(script);
        }
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Waits until specified prompt is returned by SSH server buffering response data. If expected prompt
     * is not received within defined timeout then an SshTaskTimeoutException will be thrown.
     *
     * @param command the command to send
     * @param prompt  the prompt to wait for
     * @param timeout the timeout to use when waiting for task end prompt
     *
     * @return the response from telnet server
     *
     * @throws SshException if regular expression prompts are used and regular expression is invalid
     * @see #setCommandTerminator
     * @see SshException
     * @see SshTaskTimeoutException
     */
    public synchronized String sendWait(String command, String prompt, long timeout)
            throws SshException {
        return sendWait(command, prompt, false, timeout);
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Waits until specified prompt is returned by SSH server buffering response data. If expected prompt
     * is not received within default timeout of 120 seconds then an SshTaskTimeoutException will be thrown.
     *
     * @param command the command to send
     * @param prompt  the prompt to wait for
     * @param regex   true if prompt to match is a regular expression, false otherwise
     *
     * @return the response from telnet server
     *
     * @throws SshException if regular expression prompts are used and regular expression is invalid
     * @see #setCommandTerminator
     * @see #setRegexMatchLength
     * @see SshTaskTimeoutException
     * @see SshException
     */
    public synchronized String sendWait(String command, String prompt, boolean regex)
            throws SshException {
        return sendWait(command, prompt, regex, -1);
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Waits until specified prompt is returned by SSH server buffering response data.  If expected prompt
     * is not received within default timeout of 120 seconds then an SshTaskTimeoutException will be thrown.
     *
     * @param command the command to send
     * @param prompt  the prompt to wait for
     *
     * @return the response from telnet server
     *
     * @throws SshException if regular expression prompts are used and regular expression is invalid
     * @see #setCommandTerminator
     * @see SshTaskTimeoutException
     * @see SshException
     */
    public synchronized String sendWait(String command, String prompt)
            throws SshException {
        return sendWait(command, prompt, false);
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Waits until the specified prompt is returned by SSH server buffering response data.  If shell prompt
     * is not received within default timeout of 120 seconds then an SshTaskTimeoutException will be thrown.
     *
     * @param command the command to send
     * @param prompt  the prompt to wait for
     *
     * @return the response from SSH server
     *
     * @throws SshException if regular expression prompts are used and regular expression is invalid
     * @see #setCommandTerminator
     * @see #sendWait
     * @see SshTaskTimeoutException
     * @deprecated replaced by {@link #sendWait(String, String)} method
     */
    public synchronized String send(String command, String prompt)
            throws SshException {
        return sendWait(command, prompt);
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Waits until the shell prompt is returned by SSH server buffering response data.  If shell prompt
     * is not received within defined timeout then an SshTaskTimeoutException will be thrown.
     *
     * @param command the command to send
     * @param timeout the timeout to use when waiting for shell prompt
     *
     * @return the response from SSH server
     *
     * @throws SshException if regular expression prompts are used and regular expression is invalid
     * @see #setShellPrompt
     * @see #setCommandTerminator
     * @see SshTaskTimeoutException
     */
    public synchronized String send(String command, long timeout)
            throws SshException {
        return sendWait(command, this.promptPattern, this.promptPatternIsRegex, timeout);
    }

    /**
     * Sends command to SSH host terminated by command terminator.
     * Waits until the shell prompt is returned by SSH server buffering response data.  If shell prompt
     * is not received within default timeout of 120 seconds then an SshTaskTimeoutException will be
     * thrown.
     *
     * @param command the command to send
     *
     * @return the response from SSH server
     *
     * @throws SshException if regular expression prompts are used and regular expression is invalid
     * @see #setShellPrompt
     * @see #setCommandTerminator
     * @see SshTaskTimeoutException
     */
    public synchronized String send(String command)
            throws SshException {
        return send(command, 0);
    }

    /**
     * Interrupts any running command waiting for an expected prompt.
     * This method does not stop the command from running on the server but instead stops the <code>SshSession</code>
     * from waiting for a response.  You should only use this method in the event you wish to interrupt waiting for a response.
     */
    public void interrupt() {
        this.interrupted = true;
    }

    private SshTask taskFor(String prompt, String command, String response, boolean regex, long timeout) {
        SshTask task = new SshTask(prompt, command, response);
        task.setStartPromptRegex(regex);
        task.setEndPromptRegex(regex);
        task.setRegexMatchLength(this.regexMatchLength);
        task.setTerminator(this.endOfLine);
        task.setCharacterSet(this.characterSet);

        if (timeout > 0) {
            task.setStartPromptTimeout(timeout);
            task.setEndPromptTimeout(timeout);
        }

        task.setDebug(this.ssh.getDebug());
        task.setDebugStream(ssh.getDebugStream());

        return task;
    }

    private void awaitTaskCompleted(SshTask task)
            throws SshException {
        while (!task.isComplete() && !this.interrupted && this.ssh.isConnected()) {
            Time.millis(this.pollingInterval).sleepSafe();
        }

        if (this.interrupted) {
            task.cancel();
            this.interrupted = false;
        }

        if (task.getError() != null) {
            throw new SshException(task.getError());
        }

        if (task.isTimedOut()) {
            throw new SshTaskTimeoutException(task);
        }
    }

    private void cleanup(SshScript script) {
        script.clear();
        this.ssh.removeSshListener(script);
    }

    private void debug(String message) {
        this.ssh.getDebugStream().println(message);
    }
}
