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 Guillaume Chazarain <guichaz@yahoo.fr>
  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
  29 if sys.hexversion < 0x02040000:
  30         print >> sys.stderr, 'Your python version is too old (%s)' % \
  31                                                         (sys.version.split()[0])
  32         print >> sys.stderr, 'You need at least Python 2.4'
  33         sys.exit(1)
  34
  35 from gsh import remote_dispatcher
  36 from gsh.remote_dispatcher import update_terminal_size
  37 from gsh.console import show_status, console_output
  38 from gsh import control_shell
  39 from gsh.stdin import the_stdin_thread
  40 from gsh.host_syntax import expand_syntax
  41 from gsh.version import VERSION
  42
  43 def kill_all():
  44     """When gsh quits, we kill all the remote shells we started"""
  45     for i in remote_dispatcher.all_instances():
  46         try:
  47             os.kill(i.pid, signal.SIGKILL)
  48         except OSError:
  49             # The process was already dead, no problem
  50             pass
  51
  52 def parse_cmdline():
  53     usage = '%s [OPTIONS] HOSTS...' % (sys.argv[0])
  54     parser = optparse.OptionParser(usage, version='gsh ' + VERSION)
  55     parser.add_option('--log-dir', type='str', dest='log_dir',
  56                       help='directory to log each machine conversation [none]')
  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     parser.add_option('--ssh', type='str', dest='ssh', default='ssh',
  64                       help='ssh command to use [ssh]', metavar='SSH')
  65     parser.add_option('--quick-sh', action='store_true', dest='quick_sh',
  66                       help='Do not launch a full ssh session')
  67     parser.add_option('--abort-errors', action='store_true', dest='abort_error',
  68                       help='abort if some shell fails to initialize [ignore]')
  69     parser.add_option('--debug', action='store_true', dest='debug',
  70                       help='fill the logs with debug informations')
  71     parser.add_option('--profile', action='store_true', dest='profile',
  72                       default=False, help=optparse.SUPPRESS_HELP)
  73
  74     options, args = parser.parse_args()
  75     for filename in options.hosts_filenames:
  76         try:
  77             hosts_file = open(filename, 'r')
  78             args[0:0] = [h.rstrip() for h in hosts_file.readlines()]
  79             hosts_file.close()
  80         except IOError, e:
  81             parser.error(e)
  82
  83     if not args:
  84         parser.error('no hosts given')
  85
  86     return options, args
  87
  88 def main_loop():
  89     try:
  90         while True:
  91             try:
  92                 completed, total = remote_dispatcher.count_completed_processes()
  93                 if completed == total:
  94                     # Time to use raw_input() in the stdin thread
  95                     the_stdin_thread.ready_event.set()
  96                 else:
  97                     the_stdin_thread.ready_event.clear()
  98                     # Otherwise, just print the status
  99                     show_status(completed, total)
 100                 if remote_dispatcher.all_terminated():
 101                     console_output('')
 102                     raise asyncore.ExitNow(0)
 103                 asyncore.loop(count=1, timeout=None, use_poll=True)
 104                 remote_dispatcher.handle_unfinished_lines()
 105             except KeyboardInterrupt:
 106                 control_shell.launch()
 107     except asyncore.ExitNow, e:
 108         sys.exit(e.args[0])
 109
 110 def setprocname(name):
 111     # From comments on http://davyd.livejournal.com/166352.html
 112     try:
 113         # For Python-2.5
 114         import ctypes
 115         libc = ctypes.CDLL(None)
 116         # Linux 2.6 PR_SET_NAME
 117         if libc.prctl(15, name, 0, 0, 0):
 118             # BSD
 119             libc.setproctitle(name)
 120     except:
 121         try:
 122             # For 32 bit
 123             import dl
 124             libc = dl.open(None)
 125             name += '\0'
 126             # Linux 2.6 PR_SET_NAME
 127             if libc.call('prctl', 15, name, 0, 0, 0):
 128                 # BSD
 129                 libc.call('setproctitle', name)
 130         except:
 131             pass
 132
 133 def _profile(continuation):
 134     prof_file = 'gsh.prof'
 135     try:
 136         import cProfile
 137         import pstats
 138         print 'Profiling using cProfile'
 139         cProfile.runctx('continuation()', globals(), locals(), prof_file)
 140         stats = pstats.Stats(prof_file)
 141     except ImportError:
 142         import hotshot
 143         import hotshot.stats
 144         prof = hotshot.Profile(prof_file, lineevents=1)
 145         print 'Profiling using hotshot'
 146         prof.runcall(continuation)
 147         prof.close()
 148         stats = hotshot.stats.load(prof_file)
 149     stats.strip_dirs()
 150     stats.sort_stats('time', 'calls')
 151     stats.print_stats(50)
 152     stats.print_callees(50)
 153     os.remove(prof_file)
 154
 155 def main():
 156     """Launch gsh"""
 157     locale.setlocale(locale.LC_ALL, '')
 158     setprocname('gsh')
 159     options, args = parse_cmdline()
 160     control_shell.make_singleton(options)
 161
 162     if options.log_dir:
 163         try:
 164             os.mkdir(options.log_dir)
 165         except OSError:
 166             pass # The dir already exists
 167
 168     atexit.register(kill_all)
 169     signal.signal(signal.SIGCHLD, signal.SIG_IGN) # Don't create zombies
 170     for arg in args:
 171         for host in expand_syntax(arg):
 172             remote_dispatcher.remote_dispatcher(options, host)
 173
 174     update_terminal_size()
 175     signal.signal(signal.SIGWINCH, lambda signum, frame: update_terminal_size())
 176
 177     options.interactive = not options.command and sys.stdin.isatty() and \
 178                           sys.stdout.isatty()
 179     the_stdin_thread.activate(options.interactive)
 180
 181     if options.profile:
 182         def safe_main_loop():
 183             try:
 184                 main_loop()
 185             except:
 186                 pass
 187         _profile(safe_main_loop)
 188     else:
 189         main_loop()