1#!/usr/bin/python
2
3import getopt
4import sys
5from threading import Thread
6from datetime import datetime
7import uuid
8sys.path = [".", "lib"] + sys.path
9from remote.remote_util import RemoteMachineShellConnection
10import TestInput
11
12def usage(error=None):
13    print """\
14Syntax: ssh.py [options] [command]
15
16Options:
17 -p <key=val,...> Comma-separated key=value info.
18 -i <file>        Path to .ini file containing cluster information.
19
20Available keys:
21 script=<file>           Local script to run
22 parallel=true           Run the command in parallel on all machines
23
24Examples:
25 ssh.py -i /tmp/ubuntu.ini -p script=/tmp/set_date.sh
26 ssh.py -i /tmp/ubuntu.ini -p parallel=false ls -l /tmp/core*
27"""
28    sys.exit(error)
29
30
31class CommandRunner(object):
32    def __init__(self, server, command):
33        self.server = server
34        self.command = command
35
36    def run(self):
37        remote_client = RemoteMachineShellConnection(self.server)
38        output, error = remote_client.execute_command(self.command)
39        print self.server.ip
40        print "\n".join(output)
41        print "\n".join(error)
42
43class ScriptRunner(object):
44    def __init__(self, server, script):
45        self.server = server
46        with open(script) as  f:
47            self.script_content = f.read()
48        self.script_name = "/tmp/" + str(uuid.uuid4())
49
50    def run(self):
51        remote_client = RemoteMachineShellConnection(self.server)
52        remote_client.create_file(self.script_name, self.script_content)
53        output, error = remote_client.execute_command("chmod 777 {0} ; {0} ; rm -f {0}".format(self.script_name))
54        print self.server.ip
55        print "\n".join(output)
56        print "\n".join(error)
57
58class RemoteJob(object):
59    def sequential_remote(self, input):
60        remotes = []
61        params = input.test_params
62        for server in input.servers:
63            if "script" in params:
64                remotes.append(ScriptRunner(server, params["script"]))
65            if "command" in params:
66                remotes.append(CommandRunner(server, params["command"]))
67
68        for remote in remotes:
69            try:
70                remote.run()
71            except Exception as ex:
72                print "unable to complete the job: {0}".format(ex)
73
74    def parallel_remote(self, input):
75        remotes = []
76        params = input.test_params
77        for server in input.servers:
78            if "script" in params:
79                remotes.append(ScriptRunner(server, params["script"]))
80            if "command" in params:
81                remotes.append(CommandRunner(server, params["command"]))
82
83        remote_threads = []
84        for remote in remotes:
85            remote_threads.append(Thread(target=remote.run))
86
87        for remote_thread in remote_threads:
88            remote_thread.start()
89
90        for remote_thread in remote_threads:
91            remote_thread.join()
92
93def main():
94    try:
95        (opts, args) = getopt.getopt(sys.argv[1:], 'hi:p:', [])
96        for o, a in opts:
97            if o == "-h":
98                usage()
99
100        input = TestInput.TestInputParser.get_test_input(sys.argv)
101        if not input.servers:
102            usage("ERROR: no servers specified. Please use the -i parameter.")
103    except IndexError:
104        usage()
105    except getopt.GetoptError, error:
106        usage("ERROR: " + str(error))
107
108    command_offset = 3
109    if "-p" in sys.argv[:4]:
110        command_offset += 2
111
112    command = " ".join(sys.argv[command_offset:])
113
114    if command:
115        input.test_params["command"] = command
116
117    if input.param("parallel", True):
118        # workaround for a python2.6 bug of using strptime with threads
119        datetime.strptime("30 Nov 00", "%d %b %y")
120        RemoteJob().parallel_remote(input)
121    else:
122        RemoteJob().sequential_remote(input)
123
124
125if __name__ == "__main__":
126    main()
127