gsh/main.py

   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.
   5 #
   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.
  10 #
  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.
  14 #
  15 # See the COPYING file for license information.
  16 #
  17 # Copyright (c) 2006, 2007, 2008 Guillaume Chazarain <guichaz@gmail.com>
  18
  19 # Requires python 2.4
  20
  21 import asyncore
  22 import atexit
  23 import locale
  24 import optparse
  25 import os
  26 import signal
  27 import sys
  28 import termios
  29
  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)
  35
  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
  43
  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
  52
  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)
  74
  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)
  88
  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)
  95
  96     if not args:
  97         parser.error('no hosts given')
  98
  99     return options, args
 100
 101 def find_non_interactive_command(command):
 102     if sys.stdin.isatty():
 103         return command
 104
 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
 112
 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])
 149
 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
 172
 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)
 194
 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))
 199
 200 # We handle signals in the main loop, this way we can be signaled while
 201 # handling a signal.
 202 next_signal = None
 203
 204 def main():
 205     """Launch gsh"""
 206     locale.setlocale(locale.LC_ALL, '')
 207     setprocname('gsh')
 208     options, args = parse_cmdline()
 209
 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)
 229
 230     remote_dispatcher.options = options
 231
 232     hosts = []
 233     for arg in args:
 234         hosts.extend(expand_syntax(arg))
 235
 236     dispatchers.create_remote_dispatchers(hosts)
 237
 238     signal.signal(signal.SIGWINCH, lambda signum, frame:
 239                                             dispatchers.update_terminal_size())
 240
 241     the_stdin_thread.activate(options.interactive)
 242
 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()