Show progress indicator when starting remote shells.
[polysh.git] / gsh / main.py
blob04c16b534c12061b42d30857b994cb4900f76191
1 # This program is free software; you can redistribute it and/or modify
2 # it under the terms of the GNU General Public License as published by
3 # the Free Software Foundation; either version 2 of the License, or
4 # (at your option) any later version.
6 # This program is distributed in the hope that it will be useful,
7 # but WITHOUT ANY WARRANTY; without even the implied warranty of
8 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
9 # GNU Library General Public License for more details.
11 # You should have received a copy of the GNU General Public License
12 # along with this program; if not, write to the Free Software
13 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
15 # See the COPYING file for license information.
17 # Copyright (c) 2006, 2007, 2008 Guillaume Chazarain <guichaz@gmail.com>
19 # Requires python 2.4
21 import asyncore
22 import atexit
23 import locale
24 import optparse
25 import os
26 import signal
27 import sys
28 import termios
30 if sys.hexversion < 0x02040000:
31 print >> sys.stderr, 'Your python version is too old (%s)' % \
32 (sys.version.split()[0])
33 print >> sys.stderr, 'You need at least Python 2.4'
34 sys.exit(1)
36 from gsh import remote_dispatcher
37 from gsh import dispatchers
38 from gsh.console import console_output
39 from gsh.stdin import the_stdin_thread
40 from gsh.host_syntax import expand_syntax
41 from gsh.version import VERSION
42 from gsh import control_commands
44 def kill_all():
45 """When gsh quits, we kill all the remote shells we started"""
46 for i in dispatchers.all_instances():
47 try:
48 os.kill(-i.pid, signal.SIGKILL)
49 except OSError:
50 # The process was already dead, no problem
51 pass
53 def parse_cmdline():
54 usage = '%s [OPTIONS] HOSTS...\n' % (sys.argv[0]) + \
55 'Control commands are prefixed by ":". Use :help for the list'
56 parser = optparse.OptionParser(usage, version='gsh ' + VERSION)
57 parser.add_option('--hosts-file', type='str', action='append',
58 dest='hosts_filenames', metavar='FILE', default=[],
59 help='read hostnames from given file, one per line')
60 parser.add_option('--command', type='str', dest='command', default=None,
61 help='command to execute on the remote shells',
62 metavar='CMD')
63 def_ssh = 'ssh -oLogLevel=Quiet -t %(host)s bash --noprofile --noediting'
64 parser.add_option('--ssh', type='str', dest='ssh', default=def_ssh,
65 metavar='SSH', help='ssh command to use [%s]' % def_ssh)
66 parser.add_option('--log-file', type='str', dest='log_file',
67 help='file to log each machine conversation [none]')
68 parser.add_option('--abort-errors', action='store_true', dest='abort_error',
69 help='abort if some shell fails to initialize [ignore]')
70 parser.add_option('--debug', action='store_true', dest='debug',
71 help='print debugging information.')
72 parser.add_option('--profile', action='store_true', dest='profile',
73 default=False, help=optparse.SUPPRESS_HELP)
75 options, args = parser.parse_args()
76 for filename in options.hosts_filenames:
77 try:
78 hosts_file = open(filename, 'r')
79 for line in hosts_file.readlines():
80 if '#' in line:
81 line = line[:line.index('#')]
82 line = line.strip()
83 if line:
84 args.append(line)
85 hosts_file.close()
86 except IOError, e:
87 parser.error(e)
89 if options.log_file:
90 try:
91 options.log_file = file(options.log_file, 'a')
92 except IOError, e:
93 print e
94 sys.exit(1)
96 if not args:
97 parser.error('no hosts given')
99 return options, args
101 def find_non_interactive_command(command):
102 if sys.stdin.isatty():
103 return command
105 stdin = sys.stdin.read()
106 if stdin and command:
107 print >> sys.stderr, '--command and reading from stdin are incompatible'
108 sys.exit(1)
109 if stdin and not stdin.endswith('\n'):
110 stdin += '\n'
111 return command or stdin
113 def main_loop():
114 global next_signal
115 last_status = None
116 while True:
117 try:
118 if next_signal:
119 current_signal = next_signal
120 next_signal = None
121 sig2chr = {signal.SIGINT: 'c', signal.SIGTSTP: 'z'}
122 ctrl = sig2chr[current_signal]
123 remote_dispatcher.log('> ^%c\n' % ctrl.upper())
124 control_commands.do_send_ctrl(ctrl)
125 console_output('')
126 the_stdin_thread.prepend_text = None
127 while dispatchers.count_awaited_processes()[0] and \
128 remote_dispatcher.main_loop_iteration(timeout=0.2):
129 pass
130 # Now it's quiet
131 for r in dispatchers.all_instances():
132 r.print_unfinished_line()
133 current_status = dispatchers.count_awaited_processes()
134 if current_status != last_status:
135 console_output('')
136 if remote_dispatcher.options.interactive:
137 the_stdin_thread.want_raw_input()
138 last_status = current_status
139 if dispatchers.all_terminated():
140 # Clear the prompt
141 console_output('')
142 raise asyncore.ExitNow(remote_dispatcher.options.exit_code)
143 if not next_signal:
144 # possible race here with the signal handler
145 remote_dispatcher.main_loop_iteration()
146 except asyncore.ExitNow, e:
147 console_output('')
148 sys.exit(e.args[0])
150 def setprocname(name):
151 # From comments on http://davyd.livejournal.com/166352.html
152 try:
153 # For Python-2.5
154 import ctypes
155 libc = ctypes.CDLL(None)
156 # Linux 2.6 PR_SET_NAME
157 if libc.prctl(15, name, 0, 0, 0):
158 # BSD
159 libc.setproctitle(name)
160 except:
161 try:
162 # For 32 bit
163 import dl
164 libc = dl.open(None)
165 name += '\0'
166 # Linux 2.6 PR_SET_NAME
167 if libc.call('prctl', 15, name, 0, 0, 0):
168 # BSD
169 libc.call('setproctitle', name)
170 except:
171 pass
173 def _profile(continuation):
174 prof_file = 'gsh.prof'
175 try:
176 import cProfile
177 import pstats
178 print 'Profiling using cProfile'
179 cProfile.runctx('continuation()', globals(), locals(), prof_file)
180 stats = pstats.Stats(prof_file)
181 except ImportError:
182 import hotshot
183 import hotshot.stats
184 prof = hotshot.Profile(prof_file, lineevents=1)
185 print 'Profiling using hotshot'
186 prof.runcall(continuation)
187 prof.close()
188 stats = hotshot.stats.load(prof_file)
189 stats.strip_dirs()
190 stats.sort_stats('time', 'calls')
191 stats.print_stats(50)
192 stats.print_callees(50)
193 os.remove(prof_file)
195 def restore_tty_on_exit():
196 fd = sys.stdin.fileno()
197 old = termios.tcgetattr(fd)
198 atexit.register(lambda: termios.tcsetattr(fd, termios.TCSADRAIN, old))
200 # We handle signals in the main loop, this way we can be signaled while
201 # handling a signal.
202 next_signal = None
204 def main():
205 """Launch gsh"""
206 locale.setlocale(locale.LC_ALL, '')
207 setprocname('gsh')
208 options, args = parse_cmdline()
210 atexit.register(kill_all)
211 signal.signal(signal.SIGPIPE, signal.SIG_DFL)
212 options.command = find_non_interactive_command(options.command)
213 options.exit_code = 0
214 options.interactive = not options.command and sys.stdin.isatty() and \
215 sys.stdout.isatty()
216 if options.interactive:
217 def handler(sig, frame):
218 global next_signal
219 next_signal = sig
220 signal.signal(signal.SIGINT, handler)
221 signal.signal(signal.SIGTSTP, handler)
222 restore_tty_on_exit()
223 else:
224 def handler(sig, frame):
225 signal.signal(sig, signal.SIG_DFL)
226 kill_all()
227 os.kill(0, sig)
228 signal.signal(signal.SIGINT, handler)
230 remote_dispatcher.options = options
232 hosts = []
233 for arg in args:
234 hosts.extend(expand_syntax(arg))
236 dispatchers.create_remote_dispatchers(hosts)
238 signal.signal(signal.SIGWINCH, lambda signum, frame:
239 dispatchers.update_terminal_size())
241 the_stdin_thread.activate(options.interactive)
243 if options.profile:
244 def safe_main_loop():
245 try:
246 main_loop()
247 except:
248 pass
249 _profile(safe_main_loop)
250 else:
251 main_loop()