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 os
  21 import random
  22 import zipimport
  23
  24 from gsh import callbacks
  25 from gsh import pity
  26 from gsh.console import console_output
  27 from gsh import remote_dispatcher
  28 from gsh import dispatchers
  29
  30 CMD_PREFIX = 'python -c "`echo "%s"|tr , \\\\\\n|openssl base64 -d`" '
  31
  32 CMD_SEND = CMD_PREFIX + 'send "%s" "%s" "%s"\n'
  33 CMD_FORWARD = CMD_PREFIX + 'forward "%s" "%s"\n'
  34 CMD_RECEIVE = CMD_PREFIX + 'receive "%s" "%s"\n'
  35
  36 def pity_dot_py_source():
  37     path = pity.__file__
  38     if not os.path.exists(path):
  39       try:
  40         zip_importer = zipimport.zipimporter(os.path.dirname(path))
  41       except Exception:
  42         return
  43       return zip_importer.get_source('pity')
  44     if not path.endswith('.py'):
  45         # Read from the .py source file
  46         dot_py_start = path.find('.py')
  47         if dot_py_start >= 0:
  48             path = path[:dot_py_start+3]
  49
  50     return file(path).read()
  51
  52 def base64version():
  53     python_lines = []
  54     for line in pity_dot_py_source().splitlines():
  55         hash_pos = line.find('#')
  56         if hash_pos >= 0:
  57             line = line[:hash_pos]
  58         line = line.rstrip()
  59         if line:
  60             python_lines.append(line)
  61     python_source = '\n'.join(python_lines)
  62     encoded = base64.encodestring(python_source).rstrip('\n').replace('\n', ',')
  63     return encoded
  64
  65 BASE64_PITY_PY = base64version()
  66
  67 def file_transfer_cb(dispatcher, host_port):
  68     previous_shell = get_previous_shell(dispatcher)
  69     previous_shell.dispatch_write(pity.STDIN_PREFIX + host_port + '\n')
  70
  71 def get_previous_shell(shell):
  72     shells = [i for i in dispatchers.all_instances() if i.enabled]
  73     current_pos = shells.index(shell)
  74     while True:
  75         current_pos = (current_pos - 1) % len(shells)
  76         prev_shell = shells[current_pos]
  77         if prev_shell.enabled:
  78             return prev_shell
  79
  80 def replicate(shell, path):
  81     nr_peers = len([i for i in dispatchers.all_instances() if i.enabled])
  82     if nr_peers <= 1:
  83         console_output('No other remote shell to replicate files to\n')
  84         return
  85     receiver = get_previous_shell(shell)
  86     pity_py = BASE64_PITY_PY
  87     for i in dispatchers.all_instances():
  88         if not i.enabled:
  89             continue
  90         cb = lambda host_port, i=i: file_transfer_cb(i, host_port)
  91         transfer1, transfer2 = callbacks.add('file transfer', cb, False)
  92         if i == shell:
  93             i.dispatch_command(CMD_SEND % (pity_py, path, transfer1, transfer2))
  94         elif i != receiver:
  95             i.dispatch_command(CMD_FORWARD % (pity_py, transfer1, transfer2))
  96         else:
  97             i.dispatch_command(CMD_RECEIVE % (pity_py, transfer1, transfer2))
  98         i.change_state(remote_dispatcher.STATE_RUNNING)
  99
 100 class local_uploader(remote_dispatcher.remote_dispatcher):
 101     def __init__(self, path_to_upload):
 102         remote_dispatcher.remote_dispatcher.__init__(self, '.')
 103         self.path_to_upload = path_to_upload
 104         self.upload_started = False
 105
 106     def launch_ssh(self, name):
 107         os.execl('/bin/bash', 'bash')
 108
 109     def change_state(self, state):
 110         remote_dispatcher.remote_dispatcher.change_state(self, state)
 111         if state != remote_dispatcher.STATE_IDLE:
 112             return
 113
 114         if not self.upload_started:
 115             replicate(self, self.path_to_upload)
 116             self.upload_started = True
 117         else:
 118             self.disconnect()
 119             self.close()
 120
 121 def upload(local_path):
 122     local_shell = local_uploader(local_path)
 123