polysh/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 Guillaume Chazarain <guichaz@gmail.com>
  18
  19 # Requires python 2.4
  20
  21 import asyncore
  22 import atexit
  23 import getpass
  24 import locale
  25 import optparse
  26 import os
  27 import signal
  28 import sys
  29 import termios
  30
  31 if sys.hexversion < 0x02040000:
  32         print >> sys.stderr, 'Your python version is too old (%s)' % \
  33                                                         (sys.version.split()[0])
  34         print >> sys.stderr, 'You need at least Python 2.4'
  35         sys.exit(1)
  36
  37 from polysh import remote_dispatcher
  38 from polysh import dispatchers
  39 from polysh.console import console_output
  40 from polysh.stdin import the_stdin_thread
  41 from polysh.host_syntax import expand_syntax
  42 from polysh.version import VERSION
  43 from polysh import control_commands
  44
  45 def kill_all():
  46     """When polysh quits, we kill all the remote shells we started"""
  47     for i in dispatchers.all_instances():
  48         try:
  49             os.kill(-i.pid, signal.SIGKILL)
  50         except OSError:
  51             # The process was already dead, no problem
  52             pass
  53
  54 def parse_cmdline():
  55     usage = '%s [OPTIONS] HOSTS...\n' % (sys.argv[0]) + \
  56             'Control commands are prefixed by ":". Use :help for the list'
  57     parser = optparse.OptionParser(usage, version='polysh ' + VERSION)
  58     parser.add_option('--hosts-file', type='str', action='append',
  59                       dest='hosts_filenames', metavar='FILE', default=[],
  60                       help='read hostnames from given file, one per line')
  61     parser.add_option('--command', type='str', dest='command', default=None,
  62                       help='command to execute on the remote shells',
  63                       metavar='CMD')
  64     def_ssh = 'exec ssh -oLogLevel=Quiet -t %(host)s exec bash --noprofile'
  65     parser.add_option('--ssh', type='str', dest='ssh', default=def_ssh,
  66                       metavar='SSH', help='ssh command to use [%s]' % def_ssh)
  67     parser.add_option('--user', type='str', dest='user', default=None,
  68                       help='remote user to log in as', metavar='USER')
  69     parser.add_option('--no-color', action='store_true', dest='disable_color',
  70                       help='disable colored hostnames [enabled]')
  71     parser.add_option('--password-file', type='str', dest='password_file',
  72                       default=None, metavar='FILE',
  73                       help='read a password from the specified file. - is ' +
  74                            'the tty.')
  75     parser.add_option('--log-file', type='str', dest='log_file',
  76                       help='file to log each machine conversation [none]')
  77     parser.add_option('--abort-errors', action='store_true', dest='abort_error',
  78                       help='abort if some shell fails to initialize [ignore]')
  79     parser.add_option('--debug', action='store_true', dest='debug',
  80                       help='print debugging information')
  81     parser.add_option('--profile', action='store_true', dest='profile',
  82                       default=False, help=optparse.SUPPRESS_HELP)
  83
  84     options, args = parser.parse_args()
  85     for filename in options.hosts_filenames:
  86         try:
  87             hosts_file = open(filename, 'r')
  88             for line in hosts_file.readlines():
  89                 if '#' in line:
  90                     line = line[:line.index('#')]
  91                 line = line.strip()
  92                 if line:
  93                     args.append(line)
  94             hosts_file.close()
  95         except IOError, e:
  96             parser.error(e)
  97
  98     if options.log_file:
  99         try:
 100             options.log_file = file(options.log_file, 'a')
 101         except IOError, e:
 102             print e
 103             sys.exit(1)
 104
 105     if not args:
 106         parser.error('no hosts given')
 107
 108     if options.password_file == '-':
 109         options.password = getpass.getpass()
 110     elif options.password_file is not None:
 111         password_file = file(options.password_file, 'r')
 112         options.password = password_file.readline().rstrip('\n')
 113     else:
 114         options.password = None
 115
 116     return options, args
 117
 118 def find_non_interactive_command(command):
 119     if sys.stdin.isatty():
 120         return command
 121
 122     stdin = sys.stdin.read()
 123     if stdin and command:
 124         print >> sys.stderr, '--command and reading from stdin are incompatible'
 125         sys.exit(1)
 126     if stdin and not stdin.endswith('\n'):
 127         stdin += '\n'
 128     return command or stdin
 129
 130 def main_loop():
 131     global next_signal
 132     last_status = None
 133     while True:
 134         try:
 135             if next_signal:
 136                 current_signal = next_signal
 137                 next_signal = None
 138                 sig2chr = {signal.SIGINT: 'c', signal.SIGTSTP: 'z'}
 139                 ctrl = sig2chr[current_signal]
 140                 remote_dispatcher.log('> ^%c\n' % ctrl.upper())
 141                 control_commands.do_send_ctrl(ctrl)
 142                 console_output('')
 143                 the_stdin_thread.prepend_text = None
 144             while dispatchers.count_awaited_processes()[0] and \
 145                   remote_dispatcher.main_loop_iteration(timeout=0.2):
 146                 pass
 147             # Now it's quiet
 148             for r in dispatchers.all_instances():
 149                 r.print_unfinished_line()
 150             current_status = dispatchers.count_awaited_processes()
 151             if current_status != last_status:
 152                 console_output('')
 153             if remote_dispatcher.options.interactive:
 154                 the_stdin_thread.want_raw_input()
 155             last_status = current_status
 156             if dispatchers.all_terminated():
 157                 # Clear the prompt
 158                 console_output('')
 159                 raise asyncore.ExitNow(remote_dispatcher.options.exit_code)
 160             if not next_signal:
 161                 # possible race here with the signal handler
 162                 remote_dispatcher.main_loop_iteration()
 163         except asyncore.ExitNow, e:
 164             console_output('')
 165             sys.exit(e.args[0])
 166
 167 def setprocname(name):
 168     # From comments on http://davyd.livejournal.com/166352.html
 169     try:
 170         # For Python-2.5
 171         import ctypes
 172         libc = ctypes.CDLL(None)
 173         # Linux 2.6 PR_SET_NAME
 174         if libc.prctl(15, name, 0, 0, 0):
 175             # BSD
 176             libc.setproctitle(name)
 177     except:
 178         try:
 179             # For 32 bit
 180             import dl
 181             libc = dl.open(None)
 182             name += '\0'
 183             # Linux 2.6 PR_SET_NAME
 184             if libc.call('prctl', 15, name, 0, 0, 0):
 185                 # BSD
 186                 libc.call('setproctitle', name)
 187         except:
 188             pass
 189
 190 def _profile(continuation):
 191     prof_file = 'polysh.prof'
 192     try:
 193         import cProfile
 194         import pstats
 195         print 'Profiling using cProfile'
 196         cProfile.runctx('continuation()', globals(), locals(), prof_file)
 197         stats = pstats.Stats(prof_file)
 198     except ImportError:
 199         import hotshot
 200         import hotshot.stats
 201         prof = hotshot.Profile(prof_file, lineevents=1)
 202         print 'Profiling using hotshot'
 203         prof.runcall(continuation)
 204         prof.close()
 205         stats = hotshot.stats.load(prof_file)
 206     stats.strip_dirs()
 207     stats.sort_stats('time', 'calls')
 208     stats.print_stats(50)
 209     stats.print_callees(50)
 210     os.remove(prof_file)
 211
 212 def restore_tty_on_exit():
 213     fd = sys.stdin.fileno()
 214     old = termios.tcgetattr(fd)
 215     atexit.register(lambda: termios.tcsetattr(fd, termios.TCSADRAIN, old))
 216
 217 # We handle signals in the main loop, this way we can be signaled while
 218 # handling a signal.
 219 next_signal = None
 220
 221 def main():
 222     """Launch polysh"""
 223     locale.setlocale(locale.LC_ALL, '')
 224     setprocname('polysh')
 225     options, args = parse_cmdline()
 226
 227     atexit.register(kill_all)
 228     signal.signal(signal.SIGPIPE, signal.SIG_DFL)
 229     options.command = find_non_interactive_command(options.command)
 230     options.exit_code = 0
 231     options.interactive = not options.command and sys.stdin.isatty() and \
 232                           sys.stdout.isatty()
 233     if options.interactive:
 234         def handler(sig, frame):
 235             global next_signal
 236             next_signal = sig
 237         signal.signal(signal.SIGINT, handler)
 238         signal.signal(signal.SIGTSTP, handler)
 239         restore_tty_on_exit()
 240     else:
 241       def handler(sig, frame):
 242         signal.signal(sig, signal.SIG_DFL)
 243         kill_all()
 244         os.kill(0, sig)
 245       signal.signal(signal.SIGINT, handler)
 246
 247     remote_dispatcher.options = options
 248
 249     hosts = []
 250     for arg in args:
 251         hosts.extend(expand_syntax(arg))
 252
 253     dispatchers.create_remote_dispatchers(hosts)
 254
 255     signal.signal(signal.SIGWINCH, lambda signum, frame:
 256                                             dispatchers.update_terminal_size())
 257
 258     the_stdin_thread.activate(options.interactive)
 259
 260     if options.profile:
 261         def safe_main_loop():
 262             try:
 263                 main_loop()
 264             except:
 265                 pass
 266         _profile(safe_main_loop)
 267     else:
 268         main_loop()