gsh/remote_dispatcher.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 import asyncore
  20 import os
  21 import pty
  22 import signal
  23 import sys
  24 import termios
  25
  26 from gsh.buffered_dispatcher import buffered_dispatcher
  27 from gsh import callbacks
  28 from gsh.console import console_output
  29 from gsh import display_names
  30
  31 # Either the remote shell is expecting a command or one is already running
  32 STATE_NAMES = ['not_started', 'idle', 'running', 'terminated', 'dead']
  33
  34 STATE_NOT_STARTED,         \
  35 STATE_IDLE,                \
  36 STATE_RUNNING,             \
  37 STATE_TERMINATED,          \
  38 STATE_DEAD = range(len(STATE_NAMES))
  39
  40 # Terminal color codes
  41 COLORS = [1] + range(30, 37)
  42
  43 # Count the total number of remote_dispatcher.handle_read() invocations
  44 nr_handle_read = 0
  45
  46 def main_loop_iteration(timeout=None):
  47     """Return the number of remote_dispatcher.handle_read() calls made by this
  48     iteration"""
  49     prev_nr_read = nr_handle_read
  50     asyncore.loop(count=1, timeout=timeout, use_poll=True)
  51     return nr_handle_read - prev_nr_read
  52
  53 def log(msg):
  54     if options.log_file:
  55         fd = options.log_file.fileno()
  56         while msg:
  57             try:
  58                 written = os.write(fd, msg)
  59             except OSError, e:
  60                 print 'Exception while writing log:', options.log_file.name
  61                 print e
  62                 raise asyncore.ExitNow(1)
  63             msg = msg[written:]
  64
  65 class remote_dispatcher(buffered_dispatcher):
  66     """A remote_dispatcher is a ssh process we communicate with"""
  67
  68     def __init__(self, hostname):
  69         self.pid, fd = pty.fork()
  70         if self.pid == 0:
  71             # Child
  72             self.launch_ssh(hostname)
  73             sys.exit(1)
  74
  75         # Parent
  76         buffered_dispatcher.__init__(self, fd)
  77         self.hostname = hostname
  78         self.debug = options.debug
  79         self.enabled = True # shells can be enabled and disabled
  80         self.state = STATE_NOT_STARTED
  81         self.term_size = (-1, -1)
  82         self.display_name = None
  83         self.change_name(hostname)
  84         self.init_string = self.configure_tty() + self.set_prompt()
  85         self.init_string_sent = False
  86         self.read_in_state_not_started = ''
  87         self.command = options.command
  88         self.last_printed_line = ''
  89         if sys.stdout.isatty():
  90             COLORS.insert(0, COLORS.pop()) # Rotate the colors
  91             self.color_code = COLORS[0]
  92         else:
  93             self.color_code = None
  94
  95     def launch_ssh(self, name):
  96         """Launch the ssh command in the child process"""
  97         if options.user:
  98             name = '%s@%s' % (options.user, name)
  99         evaluated = options.ssh % {'host': name}
 100         if evaluated == options.ssh:
 101             evaluated = '%s %s' % (evaluated, name)
 102         os.execlp('/bin/sh', 'sh', '-c', evaluated)
 103
 104     def set_enabled(self, enabled):
 105         if enabled != self.enabled and options.interactive:
 106             # In non-interactive mode, remote processes leave as soon
 107             # as they are terminated, but we don't want to break the
 108             # indentation if all the remaining processes have short names.
 109             display_names.set_enabled(self.display_name, enabled)
 110         self.enabled = enabled
 111
 112     def change_state(self, state):
 113         """Change the state of the remote process, logging the change"""
 114         if state is not self.state:
 115             if self.debug:
 116                 self.print_debug('state => %s' % (STATE_NAMES[state]))
 117             if self.state is STATE_NOT_STARTED:
 118                 self.read_in_state_not_started = ''
 119             self.state = state
 120
 121     def disconnect(self):
 122         """We are no more interested in this remote process"""
 123         try:
 124             os.kill(-self.pid, signal.SIGKILL)
 125         except OSError:
 126             # The process was already dead, no problem
 127             pass
 128         self.read_buffer = ''
 129         self.write_buffer = ''
 130         self.set_enabled(False)
 131         if self.read_in_state_not_started:
 132             self.print_lines(self.read_in_state_not_started)
 133             self.read_in_state_not_started = ''
 134         if options.abort_error and self.state is STATE_NOT_STARTED:
 135             raise asyncore.ExitNow(1)
 136         self.change_state(STATE_DEAD)
 137
 138     def configure_tty(self):
 139         """We don't want \n to be replaced with \r\n, and we disable the echo"""
 140         attr = termios.tcgetattr(self.fd)
 141         attr[1] &= ~termios.ONLCR # oflag
 142         attr[3] &= ~termios.ECHO # lflag
 143         termios.tcsetattr(self.fd, termios.TCSANOW, attr)
 144         # unsetopt zle prevents Zsh from resetting the tty
 145         return 'unsetopt zle 2> /dev/null;stty -echo -onlcr;'
 146
 147     def seen_prompt_cb(self, unused):
 148         if options.interactive:
 149             self.change_state(STATE_IDLE)
 150         elif self.command:
 151             p1, p2 = callbacks.add('real prompt ends', lambda d: None, True)
 152             self.dispatch_command('PS1="%s""%s\n"\n' % (p1, p2))
 153             self.dispatch_command(self.command + '\n')
 154             self.dispatch_command('exit 2>/dev/null\n')
 155             self.command = None
 156
 157     def set_prompt(self):
 158         """The prompt is important because we detect the readyness of a process
 159         by waiting for its prompt."""
 160         # No right prompt
 161         command_line = 'PS2=;RPS1=;RPROMPT=;'
 162         command_line += 'PROMPT_COMMAND=;'
 163         command_line += 'TERM=ansi;'
 164         command_line += 'unset HISTFILE;'
 165         prompt1, prompt2 = callbacks.add('prompt', self.seen_prompt_cb, True)
 166         command_line += 'PS1="%s""%s\n"\n' % (prompt1, prompt2)
 167         return command_line
 168
 169     def readable(self):
 170         """We are always interested in reading from active remote processes if
 171         the buffer is OK"""
 172         return self.state != STATE_DEAD and buffered_dispatcher.readable(self)
 173
 174     def handle_expt(self):
 175         pid, status = os.waitpid(self.pid, 0)
 176         exit_code = os.WEXITSTATUS(status)
 177         options.exit_code = max(options.exit_code, exit_code)
 178         if exit_code and options.interactive:
 179             console_output('Error talking to %s\n' % self.display_name)
 180         self.disconnect()
 181
 182     def handle_close(self):
 183         self.handle_expt()
 184
 185     def print_lines(self, lines):
 186         from gsh.display_names import max_display_name_length
 187         lines = lines.strip('\n')
 188         while True:
 189             no_empty_lines = lines.replace('\n\n', '\n')
 190             if len(no_empty_lines) == len(lines):
 191                 break
 192             lines = no_empty_lines
 193         if not lines:
 194             return
 195         indent = max_display_name_length - len(self.display_name)
 196         log_prefix = self.display_name + indent * ' ' + ' : '
 197         if self.color_code is None:
 198             console_prefix = log_prefix
 199         else:
 200             console_prefix = '\033[1;%dm%s\033[1;m' % (self.color_code,
 201                                                        log_prefix)
 202         console_data = (console_prefix +
 203                         lines.replace('\n', '\n' + console_prefix) + '\n')
 204         log_data = log_prefix + lines.replace('\n', '\n' + log_prefix) + '\n'
 205         console_output(console_data, logging_msg=log_data)
 206         self.last_printed_line = lines[lines.rfind('\n') + 1:]
 207
 208     def handle_read_fast_case(self, data):
 209         """If we are in a fast case we'll avoid the long processing of each
 210         line"""
 211         if self.state is not STATE_RUNNING or callbacks.any_in(data):
 212             # Slow case :-(
 213             return False
 214
 215         last_nl = data.rfind('\n')
 216         if last_nl == -1:
 217             # No '\n' in data => slow case
 218             return False
 219         self.read_buffer = data[last_nl + 1:]
 220         self.print_lines(data[:last_nl])
 221         return True
 222
 223     def handle_read(self):
 224         """We got some output from a remote shell, this is one of the state
 225         machine"""
 226         if self.state == STATE_DEAD:
 227             return
 228         global nr_handle_read
 229         nr_handle_read += 1
 230         new_data = buffered_dispatcher.handle_read(self)
 231         if self.debug:
 232             self.print_debug('==> ' + new_data)
 233         if self.handle_read_fast_case(self.read_buffer):
 234             return
 235         lf_pos = new_data.find('\n')
 236         if lf_pos >= 0:
 237             # Optimization: we knew there were no '\n' in the previous read
 238             # buffer, so we searched only in the new_data and we offset the
 239             # found index by the length of the previous buffer
 240             lf_pos += len(self.read_buffer) - len(new_data)
 241         elif self.state is STATE_NOT_STARTED and \
 242              options.password is not None and \
 243              'password:' in self.read_buffer.lower():
 244             self.dispatch_write(options.password + '\n')
 245             self.read_buffer = ''
 246             return
 247         while lf_pos >= 0:
 248             # For each line in the buffer
 249             line = self.read_buffer[:lf_pos + 1]
 250             if callbacks.process(line):
 251                 pass
 252             elif self.state in (STATE_IDLE, STATE_RUNNING):
 253                 self.print_lines(line)
 254             elif self.state is STATE_NOT_STARTED:
 255                 self.read_in_state_not_started += line
 256                 if 'The authenticity of host' in line:
 257                     msg = line.strip('\n') + ' Closing connection.'
 258                     self.disconnect()
 259                 elif 'REMOTE HOST IDENTIFICATION HAS CHANGED' in line:
 260                     msg = 'Remote host identification has changed.'
 261                 else:
 262                     msg = None
 263
 264                 if msg:
 265                     self.print_lines(msg + ' Consider manually connecting or ' +
 266                                      'using ssh-keyscan.')
 267
 268             # Go to the next line in the buffer
 269             self.read_buffer = self.read_buffer[lf_pos + 1:]
 270             if self.handle_read_fast_case(self.read_buffer):
 271                 return
 272             lf_pos = self.read_buffer.find('\n')
 273         if self.state is STATE_NOT_STARTED and not self.init_string_sent:
 274             self.dispatch_write(self.init_string)
 275             self.init_string_sent = True
 276
 277     def print_unfinished_line(self):
 278         """The unfinished line stayed long enough in the buffer to be printed"""
 279         if self.state is STATE_RUNNING:
 280             if not callbacks.process(self.read_buffer):
 281                 self.print_lines(self.read_buffer)
 282             self.read_buffer = ''
 283
 284     def writable(self):
 285         """Do we want to write something?"""
 286         return self.state != STATE_DEAD and buffered_dispatcher.writable(self)
 287
 288     def handle_write(self):
 289         """Let's write as much as we can"""
 290         num_sent = self.send(self.write_buffer)
 291         if self.debug:
 292             if self.state is not STATE_NOT_STARTED or options.password is None:
 293                 self.print_debug('<== ' + self.write_buffer[:num_sent])
 294         self.write_buffer = self.write_buffer[num_sent:]
 295
 296     def print_debug(self, msg):
 297         """Log some debugging information to the console"""
 298         state = STATE_NAMES[self.state]
 299         msg = msg.encode('string_escape')
 300         console_output('[dbg] %s[%s]: %s\n' % (self.display_name, state, msg))
 301
 302     def get_info(self):
 303         """Return a list with all information available about this process"""
 304         return [self.display_name, self.enabled and 'enabled' or 'disabled',
 305                 STATE_NAMES[self.state] + ':', self.last_printed_line.strip()]
 306
 307     def dispatch_write(self, buf):
 308         """There is new stuff to write when possible"""
 309         if self.state != STATE_DEAD and self.enabled:
 310             buffered_dispatcher.dispatch_write(self, buf)
 311             return True
 312
 313     def dispatch_command(self, command):
 314         if self.dispatch_write(command):
 315             self.change_state(STATE_RUNNING)
 316
 317     def change_name(self, name):
 318         """Change the name of the shell, possibly updating the maximum name
 319         length"""
 320         if not name:
 321             name = self.hostname
 322         self.display_name = display_names.change(self.display_name, name)
 323
 324     def rename(self, string):
 325         """Send to the remote shell, its new name to be shell expanded"""
 326         if string:
 327             rename1, rename2 = callbacks.add('rename', self.change_name, False)
 328             self.dispatch_command('/bin/echo "%s""%s"%s\n' %
 329                                                      (rename1, rename2, string))
 330         else:
 331             self.change_name(self.hostname)
 332
 333     def close(self):
 334         display_names.change(self.display_name, None)
 335         buffered_dispatcher.close(self)