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