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('--password-file', type='str', dest='password_file',
  70                       default=None, metavar='FILE',
  71                       help='read a password from the specified file. - is ' +
  72                            'the tty.')
  73     parser.add_option('--log-file', type='str', dest='log_file',
  74                       help='file to log each machine conversation [none]')
  75     parser.add_option('--abort-errors', action='store_true', dest='abort_error',
  76                       help='abort if some shell fails to initialize [ignore]')
  77     parser.add_option('--debug', action='store_true', dest='debug',
  78                       help='print debugging information.')
  79     parser.add_option('--profile', action='store_true', dest='profile',
  80                       default=False, help=optparse.SUPPRESS_HELP)
  81
  82     options, args = parser.parse_args()
  83     for filename in options.hosts_filenames:
  84         try:
  85             hosts_file = open(filename, 'r')
  86             for line in hosts_file.readlines():
  87                 if '#' in line:
  88                     line = line[:line.index('#')]
  89                 line = line.strip()
  90                 if line:
  91                     args.append(line)
  92             hosts_file.close()
  93         except IOError, e:
  94             parser.error(e)
  95
  96     if options.log_file:
  97         try:
  98             options.log_file = file(options.log_file, 'a')
  99         except IOError, e:
 100             print e
 101             sys.exit(1)
 102
 103     if not args:
 104         parser.error('no hosts given')
 105
 106     if options.password_file == '-':
 107         options.password = getpass.getpass()
 108     elif options.password_file is not None:
 109         password_file = file(options.password_file, 'r')
 110         options.password = password_file.readline().rstrip('\n')
 111     else:
 112         options.password = None
 113
 114     return options, args
 115
 116 def find_non_interactive_command(command):
 117     if sys.stdin.isatty():
 118         return command
 119
 120     stdin = sys.stdin.read()
 121     if stdin and command:
 122         print >> sys.stderr, '--command and reading from stdin are incompatible'
 123         sys.exit(1)
 124     if stdin and not stdin.endswith('\n'):
 125         stdin += '\n'
 126     return command or stdin
 127
 128 def main_loop():
 129     global next_signal
 130     last_status = None
 131     while True:
 132         try:
 133             if next_signal:
 134                 current_signal = next_signal
 135                 next_signal = None
 136                 sig2chr = {signal.SIGINT: 'c', signal.SIGTSTP: 'z'}
 137                 ctrl = sig2chr[current_signal]
 138                 remote_dispatcher.log('> ^%c\n' % ctrl.upper())
 139                 control_commands.do_send_ctrl(ctrl)
 140                 console_output('')
 141                 the_stdin_thread.prepend_text = None
 142             while dispatchers.count_awaited_processes()[0] and \
 143                   remote_dispatcher.main_loop_iteration(timeout=0.2):
 144                 pass
 145             # Now it's quiet
 146             for r in dispatchers.all_instances():
 147                 r.print_unfinished_line()
 148             current_status = dispatchers.count_awaited_processes()
 149             if current_status != last_status:
 150                 console_output('')
 151             if remote_dispatcher.options.interactive:
 152                 the_stdin_thread.want_raw_input()
 153             last_status = current_status
 154             if dispatchers.all_terminated():
 155                 # Clear the prompt
 156                 console_output('')
 157                 raise asyncore.ExitNow(remote_dispatcher.options.exit_code)
 158             if not next_signal:
 159                 # possible race here with the signal handler
 160                 remote_dispatcher.main_loop_iteration()
 161         except asyncore.ExitNow, e:
 162             console_output('')
 163             sys.exit(e.args[0])
 164
 165 def setprocname(name):
 166     # From comments on http://davyd.livejournal.com/166352.html
 167     try:
 168         # For Python-2.5
 169         import ctypes
 170         libc = ctypes.CDLL(None)
 171         # Linux 2.6 PR_SET_NAME
 172         if libc.prctl(15, name, 0, 0, 0):
 173             # BSD
 174             libc.setproctitle(name)
 175     except:
 176         try:
 177             # For 32 bit
 178             import dl
 179             libc = dl.open(None)
 180             name += '\0'
 181             # Linux 2.6 PR_SET_NAME
 182             if libc.call('prctl', 15, name, 0, 0, 0):
 183                 # BSD
 184                 libc.call('setproctitle', name)
 185         except:
 186             pass
 187
 188 def _profile(continuation):
 189     prof_file = 'polysh.prof'
 190     try:
 191         import cProfile
 192         import pstats
 193         print 'Profiling using cProfile'
 194         cProfile.runctx('continuation()', globals(), locals(), prof_file)
 195         stats = pstats.Stats(prof_file)
 196     except ImportError:
 197         import hotshot
 198         import hotshot.stats
 199         prof = hotshot.Profile(prof_file, lineevents=1)
 200         print 'Profiling using hotshot'
 201         prof.runcall(continuation)
 202         prof.close()
 203         stats = hotshot.stats.load(prof_file)
 204     stats.strip_dirs()
 205     stats.sort_stats('time', 'calls')
 206     stats.print_stats(50)
 207     stats.print_callees(50)
 208     os.remove(prof_file)
 209
 210 def restore_tty_on_exit():
 211     fd = sys.stdin.fileno()
 212     old = termios.tcgetattr(fd)
 213     atexit.register(lambda: termios.tcsetattr(fd, termios.TCSADRAIN, old))
 214
 215 # We handle signals in the main loop, this way we can be signaled while
 216 # handling a signal.
 217 next_signal = None
 218
 219 def main():
 220     """Launch polysh"""
 221     locale.setlocale(locale.LC_ALL, '')
 222     setprocname('polysh')
 223     options, args = parse_cmdline()
 224
 225     atexit.register(kill_all)
 226     signal.signal(signal.SIGPIPE, signal.SIG_DFL)
 227     options.command = find_non_interactive_command(options.command)
 228     options.exit_code = 0
 229     options.interactive = not options.command and sys.stdin.isatty() and \
 230                           sys.stdout.isatty()
 231     if options.interactive:
 232         def handler(sig, frame):
 233             global next_signal
 234             next_signal = sig
 235         signal.signal(signal.SIGINT, handler)
 236         signal.signal(signal.SIGTSTP, handler)
 237         restore_tty_on_exit()
 238     else:
 239       def handler(sig, frame):
 240         signal.signal(sig, signal.SIG_DFL)
 241         kill_all()
 242         os.kill(0, sig)
 243       signal.signal(signal.SIGINT, handler)
 244
 245     remote_dispatcher.options = options
 246
 247     hosts = []
 248     for arg in args:
 249         hosts.extend(expand_syntax(arg))
 250
 251     dispatchers.create_remote_dispatchers(hosts)
 252
 253     signal.signal(signal.SIGWINCH, lambda signum, frame:
 254                                             dispatchers.update_terminal_size())
 255
 256     the_stdin_thread.activate(options.interactive)
 257
 258     if options.profile:
 259         def safe_main_loop():
 260             try:
 261                 main_loop()
 262             except:
 263                 pass
 264         _profile(safe_main_loop)
 265     else:
 266         main_loop()