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