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