Web UI 自动化测试 - ssh cluster tool

package com.wbsn.mobile.util.sshserver;

import java.io.File;
import java.io.FileReader;
import java.io.IOException;

import systemobject.terminal.Prompt;

import ch.ethz.ssh2.Connection;
import ch.ethz.ssh2.SCPClient;

import com.aqua.sysobj.conn.CliCommand;
import com.aqua.sysobj.conn.CliConnectionImpl;
import com.wbsn.mobile.testcase.Constant;

public class ClusterSsh{// extends SystemObjectImpl {
    private String sshKeyLoc = "c:\\mailcontrol";
    private String tempKeyLoc;
    private String userName = "XXX";
    private String rootPass = "XXX";
    private String keyPass = "XXX";
    private String clusterHost = "";//Constant.clusterip;//"XXX.XXX.XXX.XXX";
    private String clusterPass = "t-spider";
    private String internalMachine;
    private boolean isRoot = false;
    private boolean isConnectedHost = false;
    private boolean isConnected= false;
    private CliConnectionImpl connection;    
    
    private File sshKey;
    private char[] sshKeyArray = new char[4096];
    
    public void init() throws Exception {
        //super.init();
        
    }
    
    public void close()  {
        //super.close();
        
        System.out.println("In close?");
        
        if (isConnectedHost) {
            try {
                disconnectFromHost();
            } catch (Exception e) {
                System.out.println("Unable to disconnect from remote box.");
                e.printStackTrace();
            }
        }
        
        if(connection!=null){
        
            try {
                File f = new File(this.sshKeyLoc);
                String fileName = f.getName();
                fileName = "/tmp/" + fileName;
                CliCommand command = new CliCommand("ll " + fileName);
                connection.handleCliCommand("find ssh key file", command);
                String text = command.getResult();
                if (!(text.indexOf("No such file") > -1)){
                    command = new CliCommand("rm " + tempKeyLoc);
                    connection.handleCliCommand("Remove key from firewall", command);
                }
            } catch (Exception e) {
                System.out.println("Unable to remove key from firewall.");
                e.printStackTrace();
            }
        
            connection.disconnect();
            isConnected=false;
        }
    }
    
    public void preparedEnv()throws Exception {
        
        // Add some default prompts; the firewall box and the su prompt.
        Prompt p1 = new Prompt(getPromptRe("fwl01o"),true);
        p1.setCommandEnd(true);
        
        Prompt p2 = new Prompt("Password:",false);
        p2.setCommandEnd(true);
        
        Prompt[] prompts = {p1, p2};
        
        // LinuxWbsnCliConnection is like a LinuxDefaultCliConnection except it
        // doesn't preset any prompts.
        clusterHost = Constant.clusterip;
        connection = new LinuxWbsnCliConnection(clusterHost,userName,clusterPass);
        connection.setPrompts(prompts);
        connection.init();
        
        if ("[mailcontrol@fwl01o mailcontrol]$".matches(getPromptRe("fwl01o"))) {
            System.out.println(getPromptRe("fwl01o") + " matches.");
        } else {
            System.out.println(getPromptRe("fwl01o") + " doesn't match.");
        }
        
        if (!connection.isConnected()) {
            System.out.println("Not connected.");
            //report.report("Failed to connect to cluster.",Reporter.FAIL);
            return;
        }
        
        isConnected=true;
        
        readKey(sshKeyLoc);
        tempKeyLoc = dropSshKeyWu();
        System.out.println("Ready for CLI commands.");
    }
    
    /**
     * Connects to a host within the cluster.
     *
     * @param hostName Hostname to connect to. Only include the machine name, e.g. "pmr01o".
     * @return False if already connected to a remote machine, true if successfully connected.
     * @throws Exception
     */
    public boolean connectToHost(String hostName) throws Exception {
        
        this.preparedEnv();
        
        if (isConnectedHost) {
            //report.report("Already connected to an internal machine.",Reporter.FAIL);
            return false;
        }
        
        // Add the prompt for the new machine.
        Prompt p = new Prompt(getPromptRe(hostName),true);
        p.setCommandEnd(true);
        connection.addPrompts(new Prompt[]{p});
        
        CliCommand command = new CliCommand("ssh -i " + tempKeyLoc + " " + hostName);
        //report.report("tempKeyLoc is at: "+tempKeyLoc+" hostName is: "+hostName);
        connection.handleCliCommand("SSH to box", command);
        
        //Added by Xiaobo Wang on Dec 12, 2012
        if (command.getResult().matches("(?s).*Are you sure you want to continue connecting (yes/no)*")) {
            // Enter yes
            command = new CliCommand("yes");
            connection.handleCliCommand("Enter yes", command);
        }
        
        if (!command.getResult().matches("(?s).*Enter passphrase for key.*")) {
            throw new Exception("Failed to initiate SSH connection to remote box.");
        }
        
        // Enter our password.
        command = new CliCommand(keyPass);
        connection.handleCliCommand("Enter SSH key passphrase", command);
        
        if (!command.getResult().matches("(?s)" + getPromptRe(hostName))) {
            throw new Exception("Failed to SSH to remote box.");
        }
        
        //report.report("Connected to remote box.");
        
        isConnectedHost = true;
        internalMachine = hostName;
        return true;        
    }
    
    /**
     * Disconnects from the internal host, exiting from root if needed.
     *
     * @return True if disconnected, false if not.
     * @throws Exception
     */
    public boolean disconnectFromHost() throws Exception {
        if (!isConnectedHost) {
            throw new Exception("Not connected to a box within the cluster.");
        }
        
        if (isRoot) {
            stopBeingRoot();
        }
        
        CliCommand command = new CliCommand("exit");
        connection.handleCliCommand("Closing internal connection", command);

        if (!command.getResultPrompt().getPrompt().equals(getPromptRe("fwl01o"))) {
            System.out.println("Failed to disconnect.");
            return false;
        }
        
        isConnectedHost = false;
        //report.report("Disconnected from host.");
        return true;
    }
    
    /**
     * Runs "su -" to become root.
     *
     * @throws Exception
     */
    public void becomeRoot() throws Exception {
        isRoot = true;
        
        Prompt p = new Prompt(getPromptRe(internalMachine),true);
        p.setCommandEnd(true);        
        connection.addPrompts(new Prompt[]{p});
        
        CliCommand command = new CliCommand("su -");
        connection.handleCliCommand("Become root", command);
        
        if (!command.getResultPrompt().getPrompt().equals("Password:")) {
            isRoot = false;
            throw new Exception("Was not asked for password when running su -");
        }
        
        command = new CliCommand(rootPass);
        connection.handleCliCommand("Enter root password", command);
        
        if (!command.getResultPrompt().getPrompt().equals(getPromptRe(internalMachine))) {
            isRoot = false;
            throw new Exception("Failed to become root.");
        }
        
        //report.report("Became root.");
    }
    
    /**
     * Exits from root.
     *
     * @throws Exception
     */
    public void stopBeingRoot() throws Exception {
        CliCommand command = new CliCommand("exit");
        connection.handleCliCommand("Exiting from root", command);
        isRoot = false;
        if (!command.getResultPrompt().getPrompt().equals(getPromptRe(internalMachine))) {
            isRoot = true;
            throw new Exception("Failed to exit from root.");
        }    
    }
    
    
    /**
     * Generates a regex to match the prompt for a given machine
     *
     * @param machineName The hostname of the machine
     * @return Regular expression to match the prompt
     */
    private String getPromptRe(String machineName) {
        String re = ".*\\[";
        
        if (isRoot) {
            re += "root";
        } else {
            re += userName;
        }
        
        re += "@" + machineName + ".*\\]";
        
        if (isRoot) {
            re += "#";
        } else {
            re += "\\$";
        }
        
        return re;
    }
    
    /**
     * Reads an SSH key into a char[].
     *
     * @param keyLoc The local location of the SSH key
     * @throws Exception
     */
    private void readKey(String keyLoc) throws Exception {
        System.out.println(keyLoc);
        sshKey = new File(keyLoc);
        System.out.println(sshKey.getAbsolutePath());
        FileReader reader = new FileReader(sshKey);
        if (reader.ready()) {
            try {
                reader.read(sshKeyArray);
            } catch (Exception e) {
                //report.report("Unable to read ssh key: " + e.toString(),Reporter.FAIL);
            }
        }
    }
    
    /**
     * Drops an SSH key onto the remote box. Allows for SSH within the cluster.
     *
     * @return The location of the SSH key
     */
    private String dropSshKey() throws Exception {
        CliCommand command = new CliCommand("mktemp /tmp/mailcontrol.XXXXXXXXX");
        connection.handleCliCommand("Create temp file", command);
        // We seem to get a lot of output we don't care about.
        String[] output = command.getResult().split("\n");
        String fileName = output[output.length-2].trim().replace("\n", "");

        // Drop the key on the box in a lovely way.
        command = new CliCommand("cat <<EOF > " + fileName + "\n" +
                new String(sshKeyArray) + "\n" + "EOF\n");
        connection.handleCliCommand("Drop SSH key", command);
        
        // Handle being asked for the key.
        Prompt p = new Prompt("Enter passphrase for key '" + fileName + "': ",false);
        p.setCommandEnd(true);
        connection.addPrompts(new Prompt[]{p});
                
        // And done.
        return fileName;
    }
    
 
    private String dropSshKeyWu() throws Exception {
        String fileName = "";
        File f = new File(this.sshKeyLoc);
        fileName = f.getName();
        fileName = "/tmp/" + fileName;
        CliCommand command = new CliCommand("ll " + fileName);
        connection.handleCliCommand("find ssh key file", command);
        String text = command.getResult();
        if (text.indexOf("No such file") > -1)
         this.uploadSSHKeyFile();

        Prompt p = new Prompt("Enter passphrase for key '" + fileName + "': ",
          false);
        p.setCommandEnd(true);
        connection.addPrompts(new Prompt[] { p });

        // And done.
        return fileName;
       }

    private void uploadSSHKeyFile() throws Exception {
        Connection conn = new Connection(this.clusterHost);
        conn.connect();
        File keyfile = new File(this.sshKeyLoc);
        boolean isAuthenticated = conn.authenticateWithPublicKey(this.userName,
          keyfile, this.keyPass);
        if (isAuthenticated == false)
         throw new IOException("Authentication failed.");
        SCPClient scp = conn.createSCPClient();

        scp.put(this.sshKeyLoc, "/tmp/");
        conn.close();
    }

    
    
    /**
     * Takes output from a command and removes the first and last lines, thus
     * removing the command and prompt from the output.
     *
     * @param outputToTrim The output to be trimmed.
     * @return Trimmed output.
     */
    private String getTrimmedOutput(String outputToTrim) {
        // There _MUST_ be a better way to do this, surely.
        String trimmedOutput = new String();
        
        String[] outputArray = outputToTrim.split("\\r?\\n");
        for (int i = 1; i < outputArray.length-1; i++) {
            if (i != 1) {
                trimmedOutput += "\n";
            }
            
            trimmedOutput += outputArray[i];
        }
        
        return trimmedOutput;
    }
    
    /**
     * Greps for a string in a file
     *
     * @param regex The regular expression to search for
     * @param fileName The path to the file to search
     * @return null if no matches found, otherwise all matches as reported by
     * grep
     */
    public String grep(String regex, String fileName) throws Exception {
        CliCommand command = new CliCommand("grep '" + regex + "' '" + fileName + "'");
        connection.handleCliCommand("Run grep", command);
        return getTrimmedOutput(command.getResult());
    }
    
    public String curl(String commandLine)throws Exception{
        
        CliCommand command=new CliCommand(commandLine);
        command.setTimeout(1500000);
        connection.handleCliCommand("Run curl", command);
        return getTrimmedOutput(command.getResult());
    }
    
    public String getSshKeyLoc() {
        return sshKeyLoc;
    }

    public void setSshKeyLoc(String sshKeyLoc) {
        this.sshKeyLoc = sshKeyLoc;
    }
    
    public String runCommand(String command) {
        return null;
    }

    public String getUserName() {
        return userName;
    }

    public void setUserName(String userName) {
        this.userName = userName;
    }

    public String getRootPass() {
        return rootPass;
    }

    public void setRootPass(String rootPass) {
        this.rootPass = rootPass;
    }

    public String getKeyPass() {
        return keyPass;
    }

    public void setKeyPass(String keyPass) {
        this.keyPass = keyPass;
    }

    public String getClusterHost() {
        return clusterHost;
    }

    public void setClusterHost(String clusterHost) {
        this.clusterHost = clusterHost;
    }

    public boolean isConnected() {
        return isConnected;
    }

    public void setConnected(boolean isConnected) {
        this.isConnected = isConnected;
    }

    public boolean isConnectedHost() {
        return isConnectedHost;
    }

    public void setConnectedHost(boolean isConnectedHost) {
        this.isConnectedHost = isConnectedHost;
    }



}