/*
 * SshTemplate.java
 *
 * Brazil project web application toolkit,
 * export version: 2.3 
 * Copyright (c) 2009 Sun Microsystems, Inc.
 *
 * Sun Public License Notice
 *
 * The contents of this file are subject to the Sun Public License Version 
 * 1.0 (the "License"). You may not use this file except in compliance with 
 * the License. A copy of the License is included as the file "license.terms",
 * and also available at http://www.sun.com/
 * 
 * The Original Code is from:
 *    Brazil project web application toolkit release 2.3.
 * The Initial Developer of the Original Code is: suhler.
 * Portions created by suhler are Copyright (C) Sun Microsystems, Inc.
 * All Rights Reserved.
 * 
 * Contributor(s): suhler.
 *
 * Version:  1.3
 * Created by suhler on 09/07/20
 * Last modified by suhler on 09/07/21 11:09:54
 *
 * Version Histories:
 *
 * 1.3 09/07/21-11:09:54 (suhler)
 *   Add <sshcommand> for synchronous command operation
 *
 * 1.2 09/07/20-13:18:44 (suhler)
 *   doc fixes
 *
 * 1.2 70/01/01-00:00:02 (Codemgr)
 *   SunPro Code Manager data about conflicts, renames, etc...
 *   Name history : 1 0 ssh/SshTemplate.java
 *
 * 1.1 09/07/20-13:02:51 (suhler)
 *   date and time created 09/07/20 13:02:51 by suhler
 *
 */

package sunlabs.brazil.ssh;

import java.io.OutputStream;
import java.io.IOException;
import java.util.Hashtable;
import sunlabs.brazil.server.Server;
import sunlabs.brazil.server.Request;
import sunlabs.brazil.template.RewriteContext;
import sunlabs.brazil.template.Template;
import sunlabs.brazil.util.http.HttpInputStream;
import sunlabs.brazil.util.StringMap;
import sunlabs.brazil.util.Format;
import sunlabs.brazil.template.QueueTemplate;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;

// see: http://www.ganymed.ethz.ch/ssh2/
import ch.ethz.ssh2.Session;
import ch.ethz.ssh2.Connection;

/**
 * Template to start an ssh connecton to a server.
 * @author      Stephen Uhler
 * @version	%W%
 */

public class
SshTemplate extends Template {

    Hashtable connections = new Hashtable(); // open connections
    boolean debug=false;

    /**
     * Simple version: one connection per session, no port forwarding.
     * This implements the ssh tag:
     * <pre>
     * &lt;ssh
     *   user=xx host=xx pass=xx [key=xx id=xx] stdinQ=xx stdoutQ=xx
     *   unbuffered=true|false] /&gt;
     * &lt;ssh close="x" /&gt;
     *   user=xx host=xx pass=xx [key=xxx id=xx] stdinQ=xx stdoutQ=xx
     *   unbuffered=true|false] /&gt;
     * </pre>
     * Alternately, one a connection is opened (using defer=true)
     * Commands may by run synchronously, and have their output
     * captured in the variables "stdout" and "stderr".
     * <pre>
     * &lt;ssh id=xx user=xx host=xx pass=xx [key=xx] defer=true&gt;
     * &lt;sshcommand  id=xx command=xx [stdin=xx]&gt;
     * &lt;ssh close=xx&gt;
     * </pre>
     * Tag attributes:
     * <dl>
     * <dt>host
     * <dd>The host to ssh to
     * <dt>port
     * <dd>The server port (defaults to 22)
     * <dt>user
     * <dd>The user name
     * <dt>pass
     * <dd>The user's password, if using password authentication, or
     * the password for the private key if "key" is specified.
     * <dt>key
     * <dd>The user's private key, in PEM format.  If "key" is provided, 
     * the "pass" option (if any) is used to decrypt "key". If no key
     * is provided, then password based authentication using "user"
     * and "pass" is attempted instead.
     * <dt>id
     * <dd>A unique identifier for this session.  This value will be present
     * in all Q'd output, enabling multiple sessions to share the same
     * output Q's.
     * <dt>command
     * <dd>If specified, this command is run on the remote side, otherwise
     * an interactive shell is created.
     * <dt>stdinQ
     * <dt>stdoutQ
     * <dd>The Queues (See QueueTemplate) to use to communicate with
     * the ssh streams.  The stderr/stdout of the ssh connection is sent
     * to the "stdOutQ", and is accessible using &lt;dequeue name=[stdOutQ]&gt;.
     * the output will be in the queue element "line".
     * Data can be sent to the stdin of the ssh connection with:
     * &lt;enqueue name=[stdinQ] data="line ..." /&gt;
     * <dt>unbuffered
     * <dd>by default, the "stderr" from the ssh stream is unbuffered, and
     * the "stdout" stream is line buffered.  Setting "unbuffered" to true
     * sets stderr to unbuffered as well.
     * <dt>close=[id]
     * The connection specified by "id" (if any) is closed.  When "close"
     * is specified, all other parameters are ignored.
     * </dl>
     * <p>
     * If authentication failed, the property [id].error is set.
     * When ssh terminates,  the "error" property of queue entry is set.
     * See also the example included with the distribution.
     */

    public void tag_ssh(RewriteContext hr) {
        debug(hr);
	hr.killToken();
	String host = hr.get("host");	// where to connect
	int port = Format.stringToInt(hr.get("port"), 22);
	String user = hr.get("user");
	String pass = hr.get("pass");
	String key = hr.get("key");
	String command = hr.get("command");
	String session = hr.request.props.getProperty("SessionID", "common");
	String stdoutQ = hr.get("stdoutQ", "stdoutQ." + session);
	String stdinQ = hr.get("stdinQ", "stdinQ." + session);
	boolean unbuffered=hr.isTrue("unbuffered");
	debug=hr.isTrue("debug");
	String kill = hr.get("close");

	if (kill != null) {
	    close(kill);
	    return;
	}

	if (host==null || user==null || pass==null) {
	    debug(hr, "missing parameter");
	    return;
	}
	String id = hr.get("id", user + "@" + host);

	Connection con = new Connection(host, port);
	Session s = null;
	try {
	    boolean authorized=false;
	    con.connect();
	    if (key != null && key.startsWith("-----BEGIN")) {
	        authorized = con.authenticateWithPublicKey(user,
			key.toCharArray(), pass);
	    } else {
	        authorized  = con.authenticateWithPassword(user, pass);
	    }
	    if (!authorized) {
		throw new IOException("Bad authentication");
	    }
	    connections.put(id, con);

	    // experimental Create the connection, no commands

	    if (hr.isTrue("defer")) {
		return;
	    }

	    s = con.openSession();
	    if (command != null && !command.trim().equals("")) {
		s.execCommand(command);
	    } else {
	        s.requestPTY("dumb");
	        s.startShell();
	    }
	} catch (IOException e) {
	    debug(hr, "Didn't connect: " + e);
	    hr.request.props.put(id + ".error", "Authentication failed");
	    connections.remove("id");
	    return;
	}

	if (unbuffered) {
	    new Stream(this, hr, s, stdoutQ, STDOUT_UB, id).start();
	} else {
	    new Stream(this, hr, s, stdoutQ, STDOUT, id).start();
	}
	new Stream(this, hr, s, stdoutQ, STDERR, id).start();
	new Stream(this, hr, s, stdinQ, STDIN, id).start();
    }

    /**
     * Run a command synchronously from an existing ssh connection.
     * <pre>
     * &lt;sshcommand id=xxx command=xxx [stdin=xxx timeout=sec]&gt;.
     * </pre>
     * The "Command" is run against "id" (created with &lt;ssh ...&gt;).
     * If "stdin" is supplied, it is sent to the stdin of the command.
     * The command is run to completion, and the results are placed into
     * the variables "stdout", "stderr", "status" and (sometime) error.
     * Don't forget
     * to run "&lt;ssh close=id&gt;" when done with this connection.
     */

    public void tag_sshcommand(RewriteContext hr) {
        debug(hr);
	hr.killToken();
	String id = hr.get("id");
	String command = hr.get("command");
	String stdin=hr.get("stdin");
	int timeout = Format.stringToInt(hr.get("timeout"), 0);

	if (command==null || id==null) {
	    debug(hr, "Invalid arguments");
	    return;
	}
	Connection con = (Connection) connections.get(id);
	if (con == null) {
	    debug(hr, "Unknown connection");
	    hr.request.props.put("error", "unknown connection " + id);
	    return;
	}
	Session s = null;
	try {
	    s = con.openSession();
	} catch (IOException e) {
	    System.out.println("Can't open " + e);
	    return;
	}
        Gobbler stdout = new Gobbler(s.getStdout());
        Gobbler stderr = new Gobbler(s.getStderr());
	stdout.start();
	stderr.start();

	if (stdin != null) {
	    try {
		System.out.println("setting up command stdin");
		OutputStream in = s.getStdin();
		in.write(stdin.getBytes());
		in.close();
	    } catch (IOException e) {
		System.out.println("Can't write: " + e);
	    }
	}
	try {
	    s.execCommand(command);
	    stdout.join(timeout*1000);
	    if (stdout.isAlive()) {
	        hr.request.props.put("error", "timeout");
	    }
	    stderr.join(1);
	} catch (IOException e) {
	    hr.request.props.put("error", "" + e);
	    System.out.println("io  " + e);
	} catch (InterruptedException e) {
	    System.out.println("int " + e);
	}
	hr.request.props.put("status", "" + s.getExitStatus());
	s.close();
	hr.request.props.put("stdout", stdout.get());
	hr.request.props.put("stderr", stderr.get());
    }
    
    public void close(String id) {
	Connection con = (Connection) connections.get(id);
	if (con != null) {
	    con.close();
	    connections.remove(id);
	}
    }

    static final int STDIN=0;
    static final int STDOUT=1;
    static final int STDERR=2;
    static final int STDOUT_UB=3; // unbuffered stdout
    static final int BUF_SIZE=512; // arbitrary

    static class Stream extends Thread {
	SshTemplate tmp;
	RewriteContext hr;	// used for diagnostics only
	HttpInputStream in = null;
	Request.HttpOutputStream out = null;
	String qName;		// either our input or output Queues
	String id;		// name for this command (if any)
	String type=null;	// stream type for stdout/stderr
	byte[] buf;		// buffer for stderr

	/**
	 * Start either a reader or writer.  The reader reads lines from
	 * the process, and sends them to the Queue, the writer listens from
	 * the Queue, and writes the output line to the process.
	 */

	public Stream(SshTemplate tmp, RewriteContext hr, Session ses, String qName, int what, String id) {
	    this.tmp = tmp;
	    this.hr = hr;
	    this.id = id;
	    this.qName = qName;
	    switch(what) {
	    case STDIN: // write to process stdin
	       out = new  Request.HttpOutputStream(ses.getStdin());
	       break;
	    case STDOUT: // read from process stdout
	       in = new HttpInputStream(ses.getStdout());
	       type="stdout";
	       break;
	    case STDERR:// read from process  stderr
	       in = new  HttpInputStream(ses.getStderr());
	       type="stderr";
	       buf = new byte[BUF_SIZE];
	       break;
	    case STDOUT_UB:// unbuffered stdout
	       in = new  HttpInputStream(ses.getStdout());
	       type="stdout_ub";
	       buf = new byte[BUF_SIZE];
	       break;
	    default:
	       throw new IllegalArgumentException("bad type");
	    }
	}

	public void run() {
	    // setDaemon(true);
	    boolean alive = true;
	    int count = 0;
	    while(alive) {
		if (out==null) {
		    StringMap map = new StringMap();
		    if (id != null) {
			map.add("id", id);
		    }
		    map.add("source", type);
		    map.add("count", (++count) + "");
		    String line;
		    try {
			if (!type.equals("stdout")) {
			    int n = in.read(buf, 0, buf.length);
			    if (n < 0) throw new IOException("EOF");
			    line = new String(buf, 0, n);
			} else {
			    line = in.readLine();
			}
			hr.request.log(Server.LOG_DIAGNOSTIC, hr.prefix,
				"Read: (" + line + ") Q to " + qName);
			if (line == null) {
			    throw new IOException("output closed");
			}
			map.add("line", line);
		    } catch (IOException e) {
			map.add("error", "terminated");
			alive=false;
			hr.request.log(Server.LOG_DIAGNOSTIC, hr.prefix,
				"lost output, ending");
			try {
			    tmp.close(id);
			} catch (Exception e2) {
			    System.out.println("Oops: " + e2);
			}
		    }
		    if (tmp.debug) {
			System.err.println("Enqueueing: " + map);
		    }
		    QueueTemplate.enqueue(qName, "ssh", map, false, false);
		} else {
		    QueueTemplate.QueueItem item=QueueTemplate.dequeue(qName, 10);
		    hr.request.log(Server.LOG_DIAGNOSTIC, hr.prefix,
			    "deq'd: " + item);
		    if (item != null) {
			StringMap map = (StringMap) item.data;
			try {
			    String line = map.get("line");
			    if (tmp.debug) {
				System.err.println("Sending (" + line + ")");
			    }
			    out.writeBytes(line);
			    out.flush();
			    hr.request.log(Server.LOG_DIAGNOSTIC, hr.prefix,
				"Wrote: (" + line + ")");
			} catch (IOException e) {
			    hr.request.log(Server.LOG_WARNING, hr.prefix, 
					"write failed " + e.getMessage());
			    try {out.close();} catch (Exception e2) {}
			    alive=false;
			}
		    }
		}
	    }
	}
    }

    /**
     * Gobble all the output from a stream.
     */

    static class Gobbler extends Thread {
	InputStream in;
	ByteArrayOutputStream out = new ByteArrayOutputStream();

	public Gobbler(InputStream in) {
	    this.in = in;
	}

	public String get() {
	    return out.toString();
	}

	public void run () {
            byte[] buff = new byte[BUF_SIZE];
	    while (true) {
	        try {
		    int n = in.read(buff);
		    if (n > 0) {
			out.write(buff, 0, n);
		    } else {
			break;
		    }
		} catch (IOException e) {
		    System.out.println("oops " + e);
		    break;
		}
	    }
	}
    }
}