/*
 * Copyright  2003-2005 The Apache Software Foundation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 */

package org.apache.tools.ant.taskdefs.optional.ssh;

import org.apache.tools.ant.BuildException;
import org.apache.tools.ant.Project;
import org.apache.tools.ant.util.TeeOutputStream;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.StringReader;

import java.util.Vector;

import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;

/**
 * Executes a command on a remote machine via ssh.
 *
 * @created   February 2, 2003
 * @since     Ant 1.6
 */
public class SSHExecMultiple extends SSHBase 
{
	
	private static final int BUFFER_SIZE = 1024;
	
	/** the command to execute via ssh */
	private String command = null;
	
	/** The list of commands for this session */
	private Vector commands = new Vector();
	
	/** Halt subsequent command execution on error */
	private boolean haltOnError = true;
	
	/** units are milliseconds, default is 0=infinite */
	private long maxwait = 0;
	
	/** for waiting for the command to finish */
	private Thread thread = null;
	
	private String outputProperty = null;   // like <exec>
	private File outputFile = null;   // like <exec>
	private boolean append = false;   // like <exec>
	
	private static final String TIMEOUT_MESSAGE =
		"Timeout period exceeded, connection dropped.";
	
	/**
	 * Constructor for SSHExecTask.
	 */
	public SSHExecMultiple() 
	{
		super();
	}
	
	/**
	 * Sets the command to execute on the remote host.
	 *
	 * @param command  The new command value
	 */
	public void setCommand( String command ) 
	{
		this.command = command;
	}
	
	/**
	 *  A command to execute on  the server.
	 *  A &lt;command&gt; tag was found.  Create the object,
	 *  Save it in our list, and return it.
	 */
	
	public SSHExecCommand createCommand() 
	{
		SSHExecCommand cmd = new SSHExecCommand();
		commands.addElement( cmd );
		return( cmd );
	}
	
	/** 
	 * Halt subsequent command execution on error 
	 */
	public void setHaltonerror( boolean failure ) 
	{
		haltOnError = failure;
	}
	
	public boolean getHaltonerror()
	{
        return( haltOnError );
    }
	
	/**
	 * The connection can be dropped after a specified number of
	 * milliseconds. This is sometimes useful when a connection may be
	 * flaky. Default is 0, which means &quot;wait forever&quot;.
	 *
	 * @param timeout  The new timeout value in seconds
	 */
	public void setTimeout( long timeout )  
	{
		maxwait = timeout;
	}
	
	/**
	 * If used, stores the output of the command to the given file.
	 *
	 * @param output  The file to write to.
	 */
	public void setOutput( File output ) 
	{
		outputFile = output;
	}
	
	/**
	 * Determines if the output is appended to the file given in
	 * <code>setOutput</code>. Default is false, that is, overwrite
	 * the file.
	 *
	 * @param append  True to append to an existing file, false to overwrite.
	 */
	public void setAppend( boolean append )  
	{
		this.append = append;
	}
	
	/**
	 * If set, the output of the command will be stored in the given property.
	 *
	 * @param property  The name of the property in which the command output
	 *      will be stored.
	 */
	public void setOutputproperty( String property ) 
	{
		outputProperty = property;
	}
	
	/**
	 * Execute the command on the remote host.
	 *
	 * @exception BuildException  Most likely a network error or bad parameter.
	 */
	public void execute() throws BuildException 
	{
		if( getHost() == null ) {
			throw( new BuildException( "Host is required." ) );
		}
		if( getUserInfo().getName() == null ) {
			throw( new BuildException( "Username is required." ) );
		}
		if( getUserInfo().getKeyfile() == null && getUserInfo().getPassword() == null ) {
			throw( new BuildException( "Password or Keyfile is required." ) );
		}
		if( command == null && commands.size() == 0 ) {
			throw( new BuildException( "Specify at least one command using the command attrbute or nexted <command> elements." ) );
		}
		
		Session session = null;
		
		try {
			session = openSession();
			session.setTimeout( (int)maxwait );
			
			// deal with the single command
			if( command != null ) {
				executeCommand( command, session);
			}
			
			// deal with the nested commands
			for( int i = 0; i < commands.size(); i++ ) {
				SSHExecCommand cmd = (SSHExecCommand)commands.elementAt( i );
				
				try {
					executeCommand( cmd.getCommandString(), session );
				}
				catch( BuildException e ) {
					if( getHaltonerror() ) {
						throw( e );
					} else {
						log( "Caught exception: " + e.getMessage(), Project.MSG_ERR );
					}
				}
			}
		}
		catch( BuildException e ) {
			if( getFailonerror() ) {
				throw( e );
			} else {
				log( "Caught exception: " + e.getMessage(), Project.MSG_ERR );
			}
		} 
		catch( JSchException e ) {
			if( e.getMessage().indexOf( "session is down" ) >= 0 ) {
				if( getFailonerror() ) {
					throw( new BuildException( TIMEOUT_MESSAGE, e ) );
				} else {
					log( TIMEOUT_MESSAGE, Project.MSG_ERR );
				}
			} else {
				if( getFailonerror() ) {
					throw( new BuildException( e ) );
				} else {
					log( "Caught exception: " + e.getMessage(), Project.MSG_ERR );
				}
			}
		} 
		finally {
			if( session != null && session.isConnected() ) {
				session.disconnect();
			}
		}
	}
	
	
	public void executeCommand( String cmd, Session session ) throws BuildException 
	{
		ChannelExec		savedChannel = null;
		
		try {
			
			ByteArrayOutputStream 	out 	= new ByteArrayOutputStream();
			TeeOutputStream 		tee 	= new TeeOutputStream( out, System.out );
			
			final ChannelExec 		channel = (ChannelExec)session.openChannel( "exec" );
			
			savedChannel = channel;
			
			channel.setCommand( cmd );
			channel.setOutputStream( tee );
			channel.setExtOutputStream( tee );
			channel.connect();
			
			// wait for it to finish
			thread =
				new Thread() {
				public void run() {
					while( !channel.isEOF() ) {
						if( thread == null ) {
							return;
						}
						try {
							sleep( 500 );
						} 
						catch( Exception e ) {
							// ignored
						}
					}
				}
			};
			
			thread.start();
			thread.join( maxwait );
			
			if( thread.isAlive() ) {
				// ran out of time
				thread = null;
				if( getFailonerror() || getHaltonerror() ) {
					throw( new BuildException( TIMEOUT_MESSAGE ) );
				} else {
					log( TIMEOUT_MESSAGE, Project.MSG_ERR );
				}
			} else {
				// completed successfully
				if( outputProperty != null ) {
					getProject().setProperty( outputProperty, out.toString() );
				}
				if( outputFile != null ) {
					writeToFile( out.toString(), append, outputFile );
				}
				
				// this is the wrong test if the remote OS is OpenVMS,
				// but there doesn't seem to be a way to detect it.
				int ec = channel.getExitStatus();
				if( ec != 0 ) {
					String msg = "Remote command failed with exit status " + ec;
					if( getFailonerror() || getHaltonerror() ) {
						throw( new BuildException( msg ) );
					} else {
						log( msg, Project.MSG_ERR );
					}
				}
			}
		} 
		catch( BuildException e ) {
			throw( e );
		} 
		catch( JSchException e ) {
			if( e.getMessage().indexOf( "session is down" ) >= 0 ) {
				if( getFailonerror() || getHaltonerror() ) {
					throw( new BuildException( TIMEOUT_MESSAGE, e ) );
				} else {
					log( TIMEOUT_MESSAGE, Project.MSG_ERR );
				}
			} else {
				if( getFailonerror() || getHaltonerror() ) {
					throw( new BuildException( e ) );
				} else {
					log( "Caught exception: " + e.getMessage(), Project.MSG_ERR );
				}
			}
		} 
		catch( Exception e ) {
			if( getFailonerror() || getHaltonerror() ) {
				throw( new BuildException( e ) );
			} else {
				log( "Caught exception: " + e.getMessage(), Project.MSG_ERR );
			}
		} 
		finally {
			if( savedChannel != null ) {
				savedChannel.disconnect();
			}
		}
	}
	
	
	/**
	 * Writes a string to a file. If destination file exists, it may be
	 * overwritten depending on the "append" value.
	 *
	 * @param from           string to write
	 * @param to             file to write to
	 * @param append         if true, append to existing file, else overwrite
	 * @exception Exception  most likely an IOException
	 */
	private void writeToFile( String from, boolean append, File to ) throws IOException 
	{
		FileWriter out = null;
		
		try {
			out = new FileWriter( to.getAbsolutePath(), append );
			StringReader in = new StringReader( from );
			char[] buffer = new char[ 8192 ];
			int bytesRead;
			
			while( true ) {
				bytesRead = in.read( buffer );
				if( bytesRead == -1 ) {
					break;
				}
				out.write( buffer, 0, bytesRead );
			}
			out.flush();
		} 
		finally {
			if( out != null ) {
				out.close();
			}
		}
	}
	
	/**
	 *  This class is SSHExec Command task, used to store nested command elements
	 */
	public class SSHExecCommand 
	{
		protected String commandString = "";
		
		public String getCommandString() throws BuildException 
		{
			return( commandString );
		}
		
		/**
		 *  the command as nested text
		 */
		public void addText( String s ) 
		{
			setString( getProject().replaceProperties( s ) );
		}
		
		/**
		 * the command as an attribute
		 */
		public void setString( String s ) 
		{
			commandString += s;
		}
	}
	
}