utils/src/main/java/com/cloud/utils/ssh/SSHCmdHelper.java - cloudstack - Git at Google

 //
 // 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 com.trilead.ssh2.ChannelCondition;
 import com.trilead.ssh2.Session;
 import org.apache.cloudstack.utils.security.KeyStoreUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.log4j.Logger;

 import java.io.IOException;
 import java.io.InputStream;

 public class SSHCmdHelper {
     private static final Logger s_logger = Logger.getLogger(SSHCmdHelper.class);
     private static final int DEFAULT_CONNECT_TIMEOUT = 180000;
     private static final int DEFAULT_KEX_TIMEOUT = 60000;

     public static class SSHCmdResult {
         private int returnCode = -1;
         private String stdOut;
         private String stdErr;

         public SSHCmdResult(final int returnCode, final String stdOut, final String stdErr) {
             this.returnCode = returnCode;
             this.stdOut = stdOut;
             this.stdErr = stdErr;
         }

         @Override
         public String toString() {
             return String.format("SSH cmd result: return code=%d, stdout=%s, stderr=%s",
                     getReturnCode(), getStdOut().split("-----BEGIN")[0], getStdErr());
         }

         public boolean isSuccess() {
             return returnCode == 0;
         }

         public int getReturnCode() {
             return returnCode;
         }

         public void setReturnCode(int returnCode) {
             this.returnCode = returnCode;
         }

         public String getStdOut() {
             return stdOut;
         }

         public String getStdErr() {
             return stdErr;
         }
     }

     public static com.trilead.ssh2.Connection acquireAuthorizedConnection(String ip, String username, String password) {
         return acquireAuthorizedConnection(ip, 22, username, password);
     }

     public static com.trilead.ssh2.Connection acquireAuthorizedConnection(String ip, int port, String username, String password) {
         return acquireAuthorizedConnection(ip, 22, username, password, null);
     }

     public static boolean acquireAuthorizedConnectionWithPublicKey(final com.trilead.ssh2.Connection sshConnection, final String username, final String privateKey) {
         if (StringUtils.isNotBlank(privateKey)) {
             try {
                 if (!sshConnection.authenticateWithPublicKey(username, privateKey.toCharArray(), null)) {
                     s_logger.warn("Failed to authenticate with ssh key");
                     return false;
                 }
                 return true;
             } catch (IOException e) {
                 s_logger.warn("An exception occurred when authenticate with ssh key");
                 return false;
             }
         }
         return false;
     }

     public static com.trilead.ssh2.Connection acquireAuthorizedConnection(String ip, int port, String username, String password, String privateKey) {
         com.trilead.ssh2.Connection sshConnection = new com.trilead.ssh2.Connection(ip, port);
         try {
             sshConnection.connect(null, DEFAULT_CONNECT_TIMEOUT, DEFAULT_KEX_TIMEOUT);
             if (acquireAuthorizedConnectionWithPublicKey(sshConnection, username, privateKey)) {
                 return sshConnection;
             };
             sshConnection = new com.trilead.ssh2.Connection(ip, port);
             sshConnection.connect(null, DEFAULT_CONNECT_TIMEOUT, DEFAULT_KEX_TIMEOUT);
             if (!sshConnection.authenticateWithPassword(username, password)) {
                 String[] methods = sshConnection.getRemainingAuthMethods(username);
                 StringBuffer mStr = new StringBuffer();
                 for (int i = 0; i < methods.length; i++) {
                     mStr.append(methods[i]);
                 }
                 s_logger.warn("SSH authorizes failed, support authorized methods are " + mStr);
                 return null;
             }
             return sshConnection;
         } catch (IOException e) {
             s_logger.warn("Get SSH connection failed", e);
             return null;
         }
     }

     public static void releaseSshConnection(com.trilead.ssh2.Connection sshConnection) {
         if (sshConnection != null) {
             sshConnection.close();
         }
     }

     public static boolean sshExecuteCmd(com.trilead.ssh2.Connection sshConnection, String cmd, int nTimes) {
         for (int i = 0; i < nTimes; i++) {
             try {
                 final SSHCmdResult result = sshExecuteCmdOneShot(sshConnection, cmd);
                 if (result.isSuccess()) {
                     return true;
                 }
             } catch (SshException ignored) {
                 continue;
             }
         }
         return false;
     }

     public static SSHCmdResult sshExecuteCmdWithResult(com.trilead.ssh2.Connection sshConnection, String cmd, int nTimes) {
         for (int i = 0; i < nTimes; i++) {
             try {
                 final SSHCmdResult result = sshExecuteCmdOneShot(sshConnection, cmd);
                 if (result.isSuccess()) {
                     return result;
                 }
             } catch (SshException ignored) {
                 continue;
             }
         }
         return new SSHCmdResult(-1, null, null);
     }

     public static boolean sshExecuteCmd(com.trilead.ssh2.Connection sshConnection, String cmd) {
         return sshExecuteCmd(sshConnection, cmd, 3);
     }

     public static SSHCmdResult sshExecuteCmdWithResult(com.trilead.ssh2.Connection sshConnection, String cmd) {
         return sshExecuteCmdWithResult(sshConnection, cmd, 3);
     }

     public static SSHCmdResult sshExecuteCmdOneShot(com.trilead.ssh2.Connection sshConnection, String cmd) throws SshException {
         s_logger.debug("Executing cmd: " + cmd.split(KeyStoreUtils.KS_FILENAME)[0]);
         Session sshSession = null;
         try {
             sshSession = sshConnection.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);

             if (sshSession == null) {
                 throw new SshException("Cannot open ssh session");
             }

             sshSession.execCommand(cmd);

             InputStream stdout = sshSession.getStdout();
             InputStream stderr = sshSession.getStderr();

             byte[] buffer = new byte[8192];
             StringBuffer sbStdoutResult = new StringBuffer();
             StringBuffer sbStdErrResult = new StringBuffer();

             int currentReadBytes = 0;
             while (true) {
                 if (stdout == null || stderr == null) {
                     throw new SshException("stdout or stderr of ssh session is null");
                 }
                 if ((stdout.available() == 0) && (stderr.available() == 0)) {
                     int conditions = sshSession.waitForCondition(ChannelCondition.STDOUT_DATA
                                 | ChannelCondition.STDERR_DATA | ChannelCondition.EOF | ChannelCondition.EXIT_STATUS,
                                 120000);

                     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.EXIT_STATUS) != 0) {
                         if ((conditions & (ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA)) == 0) {
                             break;
                         }
                     }

                     if ((conditions & ChannelCondition.EOF) != 0) {
                         if ((conditions & (ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA)) == 0) {
                             break;
                         }
                     }
                 }

                 while (stdout.available() > 0) {
                     currentReadBytes = stdout.read(buffer);
                     sbStdoutResult.append(new String(buffer, 0, currentReadBytes));
                 }

                 while (stderr.available() > 0) {
                     currentReadBytes = stderr.read(buffer);
                     sbStdErrResult.append(new String(buffer, 0, currentReadBytes));
                 }
             }

             final SSHCmdResult result = new SSHCmdResult(-1, sbStdoutResult.toString(), sbStdErrResult.toString());
             if (!StringUtils.isAllEmpty(result.getStdOut(), result.getStdErr())) {
                 s_logger.debug("SSH command: " + cmd.split(KeyStoreUtils.KS_FILENAME)[0] + "\nSSH command output:" + result.getStdOut().split("-----BEGIN")[0] + "\n" + result.getStdErr());
             }

             // exit status delivery might get delayed
             for(int i = 0 ; i<10 ; i++ ) {
                 Integer status = sshSession.getExitStatus();
                 if( status != null ) {
                     result.setReturnCode(status);
                     return result;
                 }
                 Thread.sleep(100);
             }
             return result;
         } catch (Exception e) {
             s_logger.debug("SSH execution failed", e);
             throw new SshException("SSH execution failed " + e.getMessage());
         } finally {
             if (sshSession != null)
                 sshSession.close();
         }
     }
 }
	//
	// 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 com.trilead.ssh2.ChannelCondition;
	import com.trilead.ssh2.Session;
	import org.apache.cloudstack.utils.security.KeyStoreUtils;
	import org.apache.commons.lang3.StringUtils;
	import org.apache.log4j.Logger;

	import java.io.IOException;
	import java.io.InputStream;

	public class SSHCmdHelper {
	private static final Logger s_logger = Logger.getLogger(SSHCmdHelper.class);
	private static final int DEFAULT_CONNECT_TIMEOUT = 180000;
	private static final int DEFAULT_KEX_TIMEOUT = 60000;

	public static class SSHCmdResult {
	private int returnCode = -1;
	private String stdOut;
	private String stdErr;

	public SSHCmdResult(final int returnCode, final String stdOut, final String stdErr) {
	this.returnCode = returnCode;
	this.stdOut = stdOut;
	this.stdErr = stdErr;
	}

	@Override
	public String toString() {
	return String.format("SSH cmd result: return code=%d, stdout=%s, stderr=%s",
	getReturnCode(), getStdOut().split("-----BEGIN")[0], getStdErr());
	}

	public boolean isSuccess() {
	return returnCode == 0;
	}

	public int getReturnCode() {
	return returnCode;
	}

	public void setReturnCode(int returnCode) {
	this.returnCode = returnCode;
	}

	public String getStdOut() {
	return stdOut;
	}

	public String getStdErr() {
	return stdErr;
	}
	}

	public static com.trilead.ssh2.Connection acquireAuthorizedConnection(String ip, String username, String password) {
	return acquireAuthorizedConnection(ip, 22, username, password);
	}

	public static com.trilead.ssh2.Connection acquireAuthorizedConnection(String ip, int port, String username, String password) {
	return acquireAuthorizedConnection(ip, 22, username, password, null);
	}

	public static boolean acquireAuthorizedConnectionWithPublicKey(final com.trilead.ssh2.Connection sshConnection, final String username, final String privateKey) {
	if (StringUtils.isNotBlank(privateKey)) {
	try {
	if (!sshConnection.authenticateWithPublicKey(username, privateKey.toCharArray(), null)) {
	s_logger.warn("Failed to authenticate with ssh key");
	return false;
	}
	return true;
	} catch (IOException e) {
	s_logger.warn("An exception occurred when authenticate with ssh key");
	return false;
	}
	}
	return false;
	}

	public static com.trilead.ssh2.Connection acquireAuthorizedConnection(String ip, int port, String username, String password, String privateKey) {
	com.trilead.ssh2.Connection sshConnection = new com.trilead.ssh2.Connection(ip, port);
	try {
	sshConnection.connect(null, DEFAULT_CONNECT_TIMEOUT, DEFAULT_KEX_TIMEOUT);
	if (acquireAuthorizedConnectionWithPublicKey(sshConnection, username, privateKey)) {
	return sshConnection;
	};
	sshConnection = new com.trilead.ssh2.Connection(ip, port);
	sshConnection.connect(null, DEFAULT_CONNECT_TIMEOUT, DEFAULT_KEX_TIMEOUT);
	if (!sshConnection.authenticateWithPassword(username, password)) {
	String[] methods = sshConnection.getRemainingAuthMethods(username);
	StringBuffer mStr = new StringBuffer();
	for (int i = 0; i < methods.length; i++) {
	mStr.append(methods[i]);
	}
	s_logger.warn("SSH authorizes failed, support authorized methods are " + mStr);
	return null;
	}
	return sshConnection;
	} catch (IOException e) {
	s_logger.warn("Get SSH connection failed", e);
	return null;
	}
	}

	public static void releaseSshConnection(com.trilead.ssh2.Connection sshConnection) {
	if (sshConnection != null) {
	sshConnection.close();
	}
	}

	public static boolean sshExecuteCmd(com.trilead.ssh2.Connection sshConnection, String cmd, int nTimes) {
	for (int i = 0; i < nTimes; i++) {
	try {
	final SSHCmdResult result = sshExecuteCmdOneShot(sshConnection, cmd);
	if (result.isSuccess()) {
	return true;
	}
	} catch (SshException ignored) {
	continue;
	}
	}
	return false;
	}

	public static SSHCmdResult sshExecuteCmdWithResult(com.trilead.ssh2.Connection sshConnection, String cmd, int nTimes) {
	for (int i = 0; i < nTimes; i++) {
	try {
	final SSHCmdResult result = sshExecuteCmdOneShot(sshConnection, cmd);
	if (result.isSuccess()) {
	return result;
	}
	} catch (SshException ignored) {
	continue;
	}
	}
	return new SSHCmdResult(-1, null, null);
	}

	public static boolean sshExecuteCmd(com.trilead.ssh2.Connection sshConnection, String cmd) {
	return sshExecuteCmd(sshConnection, cmd, 3);
	}

	public static SSHCmdResult sshExecuteCmdWithResult(com.trilead.ssh2.Connection sshConnection, String cmd) {
	return sshExecuteCmdWithResult(sshConnection, cmd, 3);
	}

	public static SSHCmdResult sshExecuteCmdOneShot(com.trilead.ssh2.Connection sshConnection, String cmd) throws SshException {
	s_logger.debug("Executing cmd: " + cmd.split(KeyStoreUtils.KS_FILENAME)[0]);
	Session sshSession = null;
	try {
	sshSession = sshConnection.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);

	if (sshSession == null) {
	throw new SshException("Cannot open ssh session");
	}

	sshSession.execCommand(cmd);

	InputStream stdout = sshSession.getStdout();
	InputStream stderr = sshSession.getStderr();

	byte[] buffer = new byte[8192];
	StringBuffer sbStdoutResult = new StringBuffer();
	StringBuffer sbStdErrResult = new StringBuffer();

	int currentReadBytes = 0;
	while (true) {
	if (stdout == null \|\| stderr == null) {
	throw new SshException("stdout or stderr of ssh session is null");
	}
	if ((stdout.available() == 0) && (stderr.available() == 0)) {
	int conditions = sshSession.waitForCondition(ChannelCondition.STDOUT_DATA
	\| ChannelCondition.STDERR_DATA \| ChannelCondition.EOF \| ChannelCondition.EXIT_STATUS,
	120000);

	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.EXIT_STATUS) != 0) {
	if ((conditions & (ChannelCondition.STDOUT_DATA \| ChannelCondition.STDERR_DATA)) == 0) {
	break;
	}
	}

	if ((conditions & ChannelCondition.EOF) != 0) {
	if ((conditions & (ChannelCondition.STDOUT_DATA \| ChannelCondition.STDERR_DATA)) == 0) {
	break;
	}
	}
	}

	while (stdout.available() > 0) {
	currentReadBytes = stdout.read(buffer);
	sbStdoutResult.append(new String(buffer, 0, currentReadBytes));
	}

	while (stderr.available() > 0) {
	currentReadBytes = stderr.read(buffer);
	sbStdErrResult.append(new String(buffer, 0, currentReadBytes));
	}
	}

	final SSHCmdResult result = new SSHCmdResult(-1, sbStdoutResult.toString(), sbStdErrResult.toString());
	if (!StringUtils.isAllEmpty(result.getStdOut(), result.getStdErr())) {
	s_logger.debug("SSH command: " + cmd.split(KeyStoreUtils.KS_FILENAME)[0] + "\nSSH command output:" + result.getStdOut().split("-----BEGIN")[0] + "\n" + result.getStdErr());
	}

	// exit status delivery might get delayed
	for(int i = 0 ; i<10 ; i++ ) {
	Integer status = sshSession.getExitStatus();
	if( status != null ) {
	result.setReturnCode(status);
	return result;
	}
	Thread.sleep(100);
	}
	return result;
	} catch (Exception e) {
	s_logger.debug("SSH execution failed", e);
	throw new SshException("SSH execution failed " + e.getMessage());
	} finally {
	if (sshSession != null)
	sshSession.close();
	}
	}
	}