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