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.
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.
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.
15 # See the COPYING file for license information.
17 # Copyright (c) 2007, 2008 Guillaume Chazarain <guichaz@gmail.com>
28 from gsh
import callbacks
30 from gsh
.console
import console_output
31 from gsh
import remote_dispatcher
32 from gsh
import dispatchers
34 def pity_dot_py_source():
36 if not os
.path
.exists(path
):
38 zip_importer
= zipimport
.zipimporter(os
.path
.dirname(path
))
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')
46 path
= path
[:dot_py_start
+3]
48 return file(path
).read()
52 for line
in pity_dot_py_source().splitlines():
53 hash_pos
= line
.find('#')
55 line
= line
[:hash_pos
]
58 python_lines
.append(line
)
59 python_source
= '\n'.join(python_lines
)
60 encoded
= base64
.encodestring(python_source
).rstrip('\n').replace('\n', ',')
63 BASE64_PITY_PY
= base64version()
65 CMD_PREFIX
= 'python -c "`echo "%s"|tr , \\\\\\n|openssl base64 -d`" ' % \
68 CMD_UPLOAD_EMIT
= ('STTY_MODE="$(stty --save)";' +
69 'stty raw &> /dev/null;' +
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'
76 def tree_max_children(depth
):
79 class file_transfer_tree_node(object):
90 self
.remote_dispatcher
= dispatcher
94 self
.is_upload
= is_upload
95 num_children
= min(len(children_dispatchers
), tree_max_children(depth
))
97 child_length
= int(math
.ceil(float(len(children_dispatchers
)) /
100 for i
in xrange(num_children
):
101 begin
= i
* child_length
102 child_dispatcher
= children_dispatchers
[begin
]
103 end
= begin
+ child_length
105 child
= file_transfer_tree_node(self
,
107 children_dispatchers
[begin
:end
],
110 self
.children
.append(child
)
111 self
.should_print_bw
= should_print_bw(self
)
112 self
.try_start_pity()
114 def host_port_cb(self
, host_port
):
115 self
.host_port
= host_port
116 self
.parent
.try_start_pity()
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
):
122 host_ports
= ' '.join(map(pipes
.quote
, host_ports
))
123 if self
.should_print_bw
:
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
)
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
)
137 cmd
= CMD_REPLICATE_EMIT
% (pipes
.quote(self
.path
), opt
, host_ports
)
138 self
.remote_dispatcher
.dispatch_command(cmd
)
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
)
149 def replicate(shell
, path
):
150 peers
= [i
for i
in dispatchers
.all_instances() if i
.enabled
]
152 console_output('No other remote shell to replicate files to\n')
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
161 sender_index
= peers
.index(shell
)
162 destinations
= peers
[:sender_index
] + peers
[sender_index
+1:]
163 tree
= file_transfer_tree_node(None,
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',
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
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)
189 os
.write(1, self
.trigger1
+ self
.trigger2
+ '\n')
190 os
._exit
(0) # The atexit handler would kill all remote shells
192 def upload_done(self
, unused
):
193 self
.first_destination
.allow_writing()
196 def upload(local_path
):
197 peers
= [i
for i
in dispatchers
.all_instances() if i
.enabled
]
199 console_output('No other remote shell to replicate files to\n')
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
208 console_output('Uploading to only one remote shell is not supported, '
212 def should_print_bw(node
, already_chosen
=[False]):
213 if not node
.children
and not already_chosen
[0]:
214 already_chosen
[0] = True
218 tree
= file_transfer_tree_node(None,