polysh/file_transfer.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) 2007 Guillaume Chazarain <guichaz@gmail.com>
  18
  19 import base64
  20 import math
  21 import os
  22 import pipes
  23 import random
  24 import subprocess
  25 import sys
  26 import zipimport
  27
  28 from polysh import callbacks
  29 from polysh import pity
  30 from polysh.console import console_output
  31 from polysh import remote_dispatcher
  32 from polysh import dispatchers
  33
  34 def pity_dot_py_source():
  35     path = pity.__file__
  36     if not os.path.exists(path):
  37       try:
  38         zip_importer = zipimport.zipimporter(os.path.dirname(path))
  39       except Exception:
  40         return
  41       return zip_importer.get_source('pity')
  42     if not path.endswith('.py'):
  43         # Read from the .py source file
  44         dot_py_start = path.find('.py')
  45         if dot_py_start >= 0:
  46             path = path[:dot_py_start+3]
  47
  48     return file(path).read()
  49
  50 def base64version():
  51     python_lines = []
  52     for line in pity_dot_py_source().splitlines():
  53         hash_pos = line.find('#')
  54         if hash_pos >= 0:
  55             line = line[:hash_pos]
  56         line = line.rstrip()
  57         if line:
  58             python_lines.append(line)
  59     python_source = '\n'.join(python_lines)
  60     encoded = base64.encodestring(python_source).rstrip('\n').replace('\n', ',')
  61     return encoded
  62
  63 def tarCreate(path):
  64     if path:
  65       path = path.rstrip('/') or '/'
  66     else:
  67       path = '.'
  68     dirname = pipes.quote(os.path.dirname(path) or '.')
  69     basename = pipes.quote(os.path.basename(path) or '/')
  70     return 'tar c -C %s %s' % (dirname, basename)
  71
  72 BASE64_PITY_PY = base64version()
  73
  74 CMD_PREFIX = 'python -c "`echo "%s"|tr , \\\\\\n|openssl base64 -d`" ' % \
  75              BASE64_PITY_PY
  76
  77 CMD_UPLOAD_EMIT = ('STTY_MODE="$(stty --save)";' +
  78                    'stty raw &> /dev/null;' +
  79                    'echo %s""%s;' +
  80                    CMD_PREFIX + ' %s upload %s;' +
  81                    'stty "$STTY_MODE"\n')
  82 CMD_REPLICATE_EMIT = '%s | ' + CMD_PREFIX + ' %s replicate %s\n'
  83 CMD_FORWARD = CMD_PREFIX + ' %s forward %s %s %s\n'
  84
  85 def tree_max_children(depth):
  86     return 2 + depth/2
  87
  88 class file_transfer_tree_node(object):
  89     def __init__(self,
  90                  parent,
  91                  dispatcher,
  92                  children_dispatchers,
  93                  depth,
  94                  should_print_bw,
  95                  path=None,
  96                  is_upload=False):
  97         self.parent = parent
  98         self.host_port = None
  99         self.remote_dispatcher = dispatcher
 100         self.children = []
 101         if path:
 102             self.path = path
 103         self.is_upload = is_upload
 104         num_children = min(len(children_dispatchers), tree_max_children(depth))
 105         if num_children:
 106             child_length = int(math.ceil(float(len(children_dispatchers)) /
 107                                          num_children))
 108             depth += 1
 109             for i in xrange(num_children):
 110                 begin = i * child_length
 111                 if begin >= len(children_dispatchers):
 112                     break
 113                 child_dispatcher = children_dispatchers[begin]
 114                 end = begin + child_length
 115                 begin += 1
 116                 child = file_transfer_tree_node(self,
 117                                                 child_dispatcher,
 118                                                 children_dispatchers[begin:end],
 119                                                 depth,
 120                                                 should_print_bw)
 121                 self.children.append(child)
 122         self.should_print_bw = should_print_bw(self)
 123         self.try_start_pity()
 124
 125     def host_port_cb(self, host_port):
 126         self.host_port = host_port
 127         self.parent.try_start_pity()
 128
 129     def try_start_pity(self):
 130         host_ports = [child.host_port for child in self.children]
 131         if len(filter(bool, host_ports)) != len(host_ports):
 132             return
 133         host_ports = ' '.join(map(pipes.quote, host_ports))
 134         if self.should_print_bw:
 135             opt = '--print-bw'
 136         else:
 137             opt = ''
 138         if self.parent:
 139             cb = lambda host_port: self.host_port_cb(host_port)
 140             t1, t2 = callbacks.add('file_transfer', cb, False)
 141             cmd = CMD_FORWARD % (opt, t1, t2, host_ports)
 142         elif self.is_upload:
 143             def start_upload(unused):
 144                 local_uploader(self.path, self.remote_dispatcher)
 145             t1, t2 = callbacks.add('upload_start', start_upload, False)
 146             cmd = CMD_UPLOAD_EMIT % (t1, t2, opt, host_ports)
 147         else:
 148             cmd = CMD_REPLICATE_EMIT % (tarCreate(self.path), opt, host_ports)
 149         self.remote_dispatcher.dispatch_command(cmd)
 150
 151     def __str__(self):
 152         children_str = ''
 153         for child in self.children:
 154             child_str = str(child)
 155             for line in child_str.splitlines():
 156                 children_str += '+--%s\n' % line
 157         return '%s\n%s' % (self.remote_dispatcher.display_name, children_str)
 158
 159
 160 def replicate(shell, path):
 161     peers = [i for i in dispatchers.all_instances() if i.enabled]
 162     if len(peers) <= 1:
 163         console_output('No other remote shell to replicate files to\n')
 164         return
 165
 166     def should_print_bw(node, already_chosen=[False]):
 167         if not node.children and not already_chosen[0] and not node.is_upload:
 168             already_chosen[0] = True
 169             return True
 170         return False
 171
 172     sender_index = peers.index(shell)
 173     destinations = peers[:sender_index] + peers[sender_index+1:]
 174     tree = file_transfer_tree_node(None,
 175                                    shell,
 176                                    destinations,
 177                                    0,
 178                                    should_print_bw,
 179                                    path=path)
 180
 181
 182 class local_uploader(remote_dispatcher.remote_dispatcher):
 183     def __init__(self, path_to_upload, first_destination):
 184         self.path_to_upload = path_to_upload
 185         self.trigger1, self.trigger2 = callbacks.add('upload_done',
 186                                                      self.upload_done,
 187                                                      False)
 188         self.first_destination = first_destination
 189         self.first_destination.drain_and_block_writing()
 190         remote_dispatcher.remote_dispatcher.__init__(self, '.')
 191         self.temporary = True
 192
 193     def launch_ssh(self, name):
 194         cmd = '%s | (openssl base64; echo %s) >&%d' % (
 195             tarCreate(self.path_to_upload),
 196             pity.BASE64_TERMINATOR,
 197             self.first_destination.fd)
 198         subprocess.call(cmd, shell=True)
 199
 200         os.write(1, self.trigger1 + self.trigger2 + '\n')
 201         os._exit(0)  # The atexit handler would kill all remote shells
 202
 203     def upload_done(self, unused):
 204         self.first_destination.allow_writing()
 205
 206
 207 def upload(local_path):
 208     peers = [i for i in dispatchers.all_instances() if i.enabled]
 209     if not peers:
 210         console_output('No other remote shell to replicate files to\n')
 211         return
 212
 213     if len(peers) == 1:
 214         # We wouldn't be able to show the progress indicator with only one
 215         # destination. We need one remote connection in blocking mode to send
 216         # the base64 data to. We also need one remote connection in non blocking
 217         # mode for polysh to display the progress indicator via the main select
 218         # loop.
 219         console_output('Uploading to only one remote shell is not supported, '
 220                        'use scp instead\n')
 221         return
 222
 223     def should_print_bw(node, already_chosen=[False]):
 224         if not node.children and not already_chosen[0]:
 225             already_chosen[0] = True
 226             return True
 227         return False
 228
 229     tree = file_transfer_tree_node(None,
 230                                    peers[0],
 231                                    peers[1:],
 232                                    0,
 233                                    should_print_bw,
 234                                    path=local_path,
 235                                    is_upload=True)