utils/src/com/cloud/utils/ssh/SshHelper.java

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 com.cloud.utils.ssh;

import java.io.File;
import java.io.InputStream;

import org.apache.log4j.Logger;

import com.cloud.utils.Pair;
import com.trilead.ssh2.ChannelCondition;

public class SshHelper {
	private static int DEFAULT_CONNECT_TIMEOUT = 60000;
	private static int DEFAULT_KEX_TIMEOUT = 60000;
	
	private static final Logger s_logger = Logger.getLogger(SshHelper.class);
	
	public static Pair<Boolean, String> sshExecute(String host, int port, String user, File pemKeyFile, String password, String command) 
		throws Exception {
	
		return sshExecute(host, port, user, pemKeyFile, password, command, 60000, 60000, 120000);
	}
	
	public static void scpTo(String host, int port, String user, File pemKeyFile, String password, String remoteTargetDirectory, 
		String localFile, String fileMode) throws Exception {
		
		scpTo(host, port, user, pemKeyFile, password, remoteTargetDirectory, localFile, fileMode, 
			DEFAULT_CONNECT_TIMEOUT, DEFAULT_KEX_TIMEOUT);
	}
	
	public static void scpTo(String host, int port, String user, File pemKeyFile, String password, String remoteTargetDirectory, 
		byte[] data, String remoteFileName, String fileMode) throws Exception {
		
		scpTo(host, port, user, pemKeyFile, password, remoteTargetDirectory, data, remoteFileName, fileMode, 
			DEFAULT_CONNECT_TIMEOUT, DEFAULT_KEX_TIMEOUT);
	}
	
	public static void scpTo(String host, int port, String user, File pemKeyFile, String password, String remoteTargetDirectory, 
		String localFile, String fileMode, int connectTimeoutInMs, int kexTimeoutInMs) throws Exception {
		
		com.trilead.ssh2.Connection conn = null;
		com.trilead.ssh2.SCPClient scpClient = null;
		
		try {
			conn = new com.trilead.ssh2.Connection(host, port);
			conn.connect(null, connectTimeoutInMs, kexTimeoutInMs);
			
			if(pemKeyFile == null) {
				if(!conn.authenticateWithPassword(user, password)) {
					String msg = "Failed to authentication SSH user " + user + " on host " + host;
					s_logger.error(msg);
					throw new Exception(msg);
				}
			} else {
				if(!conn.authenticateWithPublicKey(user, pemKeyFile, password)) {
					String msg = "Failed to authentication SSH user " + user + " on host " + host;
					s_logger.error(msg);
					throw new Exception(msg);
				}
			}
			
			scpClient = conn.createSCPClient();
			
			if(fileMode != null)
				scpClient.put(localFile, remoteTargetDirectory, fileMode);
			else
				scpClient.put(localFile, remoteTargetDirectory);
		} finally {
			if(conn != null)
				conn.close();
		}
	}
	
	public static void scpTo(String host, int port, String user, File pemKeyFile, String password, String remoteTargetDirectory, 
		byte[] data, String remoteFileName, String fileMode, int connectTimeoutInMs, int kexTimeoutInMs) throws Exception {
		
		com.trilead.ssh2.Connection conn = null;
		com.trilead.ssh2.SCPClient scpClient = null;
		
		try {
			conn = new com.trilead.ssh2.Connection(host, port);
			conn.connect(null, connectTimeoutInMs, kexTimeoutInMs);
			
			if(pemKeyFile == null) {
				if(!conn.authenticateWithPassword(user, password)) {
					String msg = "Failed to authentication SSH user " + user + " on host " + host;
					s_logger.error(msg);
					throw new Exception(msg);
				}
			} else {
				if(!conn.authenticateWithPublicKey(user, pemKeyFile, password)) {
					String msg = "Failed to authentication SSH user " + user + " on host " + host;
					s_logger.error(msg);
					throw new Exception(msg);
				}
			}
			
			scpClient = conn.createSCPClient();
			if(fileMode != null)
				scpClient.put(data, remoteFileName, remoteTargetDirectory, fileMode);
			else
				scpClient.put(data, remoteFileName, remoteTargetDirectory);
		} finally {
			if(conn != null)
				conn.close();
		}
	}
	
	public static Pair<Boolean, String> sshExecute(String host, int port, String user, File pemKeyFile, String password, String command, 
		int connectTimeoutInMs, int kexTimeoutInMs, int waitResultTimeoutInMs) throws Exception {
		
		com.trilead.ssh2.Connection conn = null;
		com.trilead.ssh2.Session sess = null; 
		try {
			conn = new com.trilead.ssh2.Connection(host, port);
			conn.connect(null, connectTimeoutInMs, kexTimeoutInMs);

			if(pemKeyFile == null) {
				if(!conn.authenticateWithPassword(user, password)) {
					String msg = "Failed to authentication SSH user " + user + " on host " + host;
					s_logger.error(msg);
					throw new Exception(msg);
				}
			} else {
				if(!conn.authenticateWithPublicKey(user, pemKeyFile, password)) {
					String msg = "Failed to authentication SSH user " + user + " on host " + host;
					s_logger.error(msg);
					throw new Exception(msg);
				}
			}
			sess = conn.openSession();
			
			// There is a bug in Trilead library, wait a second before
			// starting a shell and executing commands, from http://spci.st.ewi.tudelft.nl/chiron/xref/nl/tudelft/swerl/util/SSHConnection.html
			Thread.sleep(1000);
			
			sess.execCommand(command);
			
			InputStream stdout = sess.getStdout();
			InputStream stderr = sess.getStderr();
	
            byte[] buffer = new byte[8192];
            StringBuffer sbResult = new StringBuffer();
            
            int currentReadBytes = 0;
            while (true) {
                if ((stdout.available() == 0) && (stderr.available() == 0)) {
                    int conditions = sess.waitForCondition(ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA | ChannelCondition.EOF, 
                        waitResultTimeoutInMs);
                    
                    if ((conditions & ChannelCondition.TIMEOUT) != 0) {
                        String msg = "Timed out in waiting SSH execution result";
                        s_logger.error(msg);
                        throw new Exception(msg);
                    }

                    if ((conditions & ChannelCondition.EOF) != 0) {
                        if ((conditions & (ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA)) == 0) {                            
                            break;
                        }
                    }
                }
                            
                while (stdout.available() > 0) {
                    currentReadBytes = stdout.read(buffer);
                    sbResult.append(new String(buffer, 0, currentReadBytes));
                }
            
                while (stderr.available() > 0) {
                    currentReadBytes = stderr.read(buffer);
                    sbResult.append(new String(buffer, 0, currentReadBytes));
                }
            }
            
            String result = sbResult.toString();
            if (sess.getExitStatus() != null && sess.getExitStatus().intValue() != 0) {
                s_logger.error("SSH execution of command " + command + " has an error status code in return. result output: " + result);
                return new Pair<Boolean, String>(false, result);
            }
            
            return new Pair<Boolean, String>(true, result);
		} finally {
			if(sess != null)
				sess.close();
			
			if(conn != null)
				conn.close();
		}
	}
}