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 information')
  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             for line in hosts_file.readlines():
  79                 if '#' in line:
  80                     line = line[:line.index('#')]
  81                 line = line.strip()
  82                 if line:
  83                     args.append(line)
  84             hosts_file.close()
  85         except IOError, e:
  86             parser.error(e)
  87
  88     if not args:
  89         parser.error('no hosts given')
  90
  91     return options, args
  92
  93 def main_loop():
  94     try:
  95         while True:
  96             try:
  97                 completed, total = remote_dispatcher.count_completed_processes()
  98                 if completed == total:
  99                     # Time to use raw_input() in the stdin thread
 100                     the_stdin_thread.ready_event.set()
 101                 else:
 102                     the_stdin_thread.ready_event.clear()
 103                     # Otherwise, just print the status
 104                     show_status(completed, total)
 105                 if remote_dispatcher.all_terminated():
 106                     console_output('')
 107                     raise asyncore.ExitNow(0)
 108                 asyncore.loop(count=1, timeout=None, use_poll=True)
 109                 remote_dispatcher.handle_unfinished_lines()
 110             except KeyboardInterrupt:
 111                 control_shell.launch()
 112     except asyncore.ExitNow, e:
 113         sys.exit(e.args[0])
 114
 115 def setprocname(name):
 116     # From comments on http://davyd.livejournal.com/166352.html
 117     try:
 118         # For Python-2.5
 119         import ctypes
 120         libc = ctypes.CDLL(None)
 121         # Linux 2.6 PR_SET_NAME
 122         if libc.prctl(15, name, 0, 0, 0):
 123             # BSD
 124             libc.setproctitle(name)
 125     except:
 126         try:
 127             # For 32 bit
 128             import dl
 129             libc = dl.open(None)
 130             name += '\0'
 131             # Linux 2.6 PR_SET_NAME
 132             if libc.call('prctl', 15, name, 0, 0, 0):
 133                 # BSD
 134                 libc.call('setproctitle', name)
 135         except:
 136             pass
 137
 138 def _profile(continuation):
 139     prof_file = 'gsh.prof'
 140     try:
 141         import cProfile
 142         import pstats
 143         print 'Profiling using cProfile'
 144         cProfile.runctx('continuation()', globals(), locals(), prof_file)
 145         stats = pstats.Stats(prof_file)
 146     except ImportError:
 147         import hotshot
 148         import hotshot.stats
 149         prof = hotshot.Profile(prof_file, lineevents=1)
 150         print 'Profiling using hotshot'
 151         prof.runcall(continuation)
 152         prof.close()
 153         stats = hotshot.stats.load(prof_file)
 154     stats.strip_dirs()
 155     stats.sort_stats('time', 'calls')
 156     stats.print_stats(50)
 157     stats.print_callees(50)
 158     os.remove(prof_file)
 159
 160 def main():
 161     """Launch gsh"""
 162     locale.setlocale(locale.LC_ALL, '')
 163     setprocname('gsh')
 164     options, args = parse_cmdline()
 165     control_shell.make_singleton(options)
 166
 167     if options.log_dir:
 168         try:
 169             os.mkdir(options.log_dir)
 170         except OSError:
 171             pass # The dir already exists
 172
 173     atexit.register(kill_all)
 174     signal.signal(signal.SIGCHLD, signal.SIG_IGN) # Don't create zombies
 175     for arg in args:
 176         for host in expand_syntax(arg):
 177             remote_dispatcher.remote_dispatcher(options, host)
 178
 179     update_terminal_size()
 180     signal.signal(signal.SIGWINCH, lambda signum, frame: update_terminal_size())
 181
 182     options.interactive = not options.command and sys.stdin.isatty() and \
 183                           sys.stdout.isatty()
 184     the_stdin_thread.activate(options.interactive)
 185
 186     if options.profile:
 187         def safe_main_loop():
 188             try:
 189                 main_loop()
 190             except:
 191                 pass
 192         _profile(safe_main_loop)
 193     else:
 194         main_loop()