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