def _start_slave(src_id, cmdline, router):
"""
This runs in the target context, it is invoked by _fakessh_main running in
the fakessh context immediately after startup. It starts the slave process
(the the point where it has a stdin_handle to target but not stdout_chan to
write to), and waits for main to.
"""
LOG.debug('_start_slave(%r, %r)', router, cmdline)
proc = subprocess.Popen(
cmdline,
# SSH server always uses user's shell.
shell=True,
# SSH server always executes new commands in the user's HOME.
cwd=os.path.expanduser('~'),
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
)
process = Process(
router,
proc.stdin.fileno(),
proc.stdout.fileno(),
proc,
)
return process.control_handle, process.stdin_handle
def connect(path, broker=None):
LOG.debug('unix.connect(path=%r)', path)
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
sock.connect(path)
sock.send(struct.pack('>L', os.getpid()))
mitogen.context_id, remote_id, pid = struct.unpack('>LLL', sock.recv(12))
mitogen.parent_id = remote_id
mitogen.parent_ids = [remote_id]
LOG.debug('unix.connect(): local ID is %r, remote is %r',
mitogen.context_id, remote_id)
router = mitogen.master.Router(broker=broker)
stream = mitogen.core.Stream(router, remote_id)
stream.accept(sock.fileno(), sock.fileno())
stream.name = u'unix_listener.%d' % (pid,)
context = mitogen.parent.Context(router, remote_id)
router.register(context, stream)
mitogen.core.listen(router.broker, 'shutdown',
lambda: router.disconnect_stream(stream))
sock.close()
return router, context
def _on_get_module(self, msg):
LOG.debug('%r._on_get_module(%r)', self, msg)
if msg == mitogen.core._DEAD:
return
fullname = msg.data
callback = lambda: self._on_cache_callback(msg, fullname)
self.importer._request_module(fullname, callback)
def run_once(self):
try:
msg = self.recv.get()
except mitogen.core.ChannelError, e:
# Channel closed due to broker shutdown, exit gracefully.
LOG.debug('%r: channel closed: %s', self, e)
self.running = False
return
def store_and_forward(self, path, data, context):
LOG.debug('%r.store_and_forward(%r, %r, %r)', self, path, data,
context)
waiters = self._store(path, data)
if context.context_id != mitogen.context_id:
self._forward(context, path)
for callback in waiters:
callback()
def upgrade(self, importer, parent):
LOG.debug('%r.upgrade()', self)
self.id_allocator = ChildIdAllocator(router=self)
self.responder = ModuleForwarder(
router=self,
parent_context=parent,
importer=importer,
)
self.route_monitor = RouteMonitor(self, parent)
def shutdown(self, wait=False):
LOG.debug('%r.shutdown() sending SHUTDOWN', self)
latch = mitogen.core.Latch()
mitogen.core.listen(self, 'disconnect', lambda: latch.put(None))
self.send(mitogen.core.Message(handle=mitogen.core.SHUTDOWN, ))
if wait:
latch.get()
def on_disconnect(self, broker):
pid, status = os.waitpid(self.pid, os.WNOHANG)
if pid:
LOG.debug('%r: child process exit status was %d', self, status)
else:
LOG.debug('%r: child process still alive, sending SIGTERM', self)
os.kill(self.pid, signal.SIGTERM)
pid, status = os.waitpid(self.pid, 0)
super(Stream, self).on_disconnect(broker)
def connect(context):
"""Connect to a Broker at the address specified in our associated
Context."""
LOG.debug('%s.connect()', __name__)
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.receive_side = mitogen.core.Side(self, sock.fileno())
self.transmit_side = mitogen.core.Side(self, sock.fileno())
sock.connect(self._context.parent_addr)
self.enqueue(0, self._context.name)
def on_allocate_id(self, msg):
if msg.is_dead:
return
id_, last_id = self.allocate_block()
requestee = self.router.context_by_id(msg.src_id)
LOG.debug('%r: allocating [%r..%r) to %r',
self, id_, last_id, requestee)
msg.reply((id_, last_id))
def _find_one_component(self, modname, search_path):
try:
#fp, path, (suffix, _, kind) = imp.find_module(modname, search_path)
return imp.find_module(modname, search_path)
except ImportError:
e = sys.exc_info()[1]
LOG.debug('%r: imp.find_module(%r, %r) -> %s', self, modname,
[search_path], e)
return None
def _found_package(self, fullname, path):
path = os.path.join(path, '__init__.py')
LOG.debug('%r: %r is PKG_DIRECTORY: %r', self, fullname, path)
return self._found_module(
fullname=fullname,
path=path,
fp=open(path, 'rb'),
is_pkg=True,
)
def propagate_to(self, context, path):
LOG.debug('%r.propagate_to(%r, %r)', self, context, path)
if path not in self._cache:
fp = open(path, 'rb')
try:
self._cache[path] = mitogen.core.Blob(fp.read())
finally:
fp.close()
self._forward(context, path)
def on_shutdown(self, broker):
"""Request the slave gracefully shut itself down."""
LOG.debug('%r closing CALL_FUNCTION channel', self)
self.send(
mitogen.core.Message(
src_id=mitogen.context_id,
dst_id=self.remote_id,
handle=mitogen.core.SHUTDOWN,
))
def add_route(self, target_id, stream):
LOG.debug('%r.add_route(%r, %r)', self, target_id, stream)
assert isinstance(target_id, int)
assert isinstance(stream, Stream)
try:
self._stream_by_id[target_id] = stream
except KeyError:
LOG.error('%r: cant add route to %r via %r: no such stream',
self, target_id, stream)
def _on_get_module(self, msg):
LOG.debug('%r._on_get_module(%r)', self, msg)
if msg.is_dead:
return
fullname = msg.data
self.importer._request_module(fullname,
lambda: self._on_cache_callback(msg, fullname)
)
def _on_control(self, msg):
if not msg.is_dead:
command, arg = msg.unpickle(throw=False)
LOG.debug('%r._on_control(%r, %s)', self, command, arg)
func = getattr(self, '_on_%s' % (command,), None)
if func:
return func(msg, arg)
LOG.warning('%r: unknown command %r', self, command)
def _on_control(self, msg):
if not msg.is_dead:
command, arg = msg.unpickle(throw=False)
LOG.debug('%r._on_control(%r, %s)', self, command, arg)
func = getattr(self, '_on_%s' % (command, ), None)
if func:
return func(msg, arg)
LOG.warning('%r: unknown command %r', self, command)
def connect(self):
LOG.debug('%r.connect()', self)
pid, fd = self.create_child(*self.get_boot_command())
self.name = 'local.%s' % (pid, )
self.receive_side = mitogen.core.Side(self, fd)
self.transmit_side = mitogen.core.Side(self, os.dup(fd))
LOG.debug('%r.connect(): child process stdin/stdout=%r', self,
self.receive_side.fd)
self._connect_bootstrap()
def _on_control(self, msg):
if msg != mitogen.core._DEAD:
command, arg = msg.unpickle()
LOG.debug('%r._on_control(%r, %s)', self, command, arg)
func = getattr(self, '_on_%s' % (command, ), None)
if func:
return func(msg, arg)
LOG.warning('%r: unknown command %r', self, command)
def _on_detaching(self, msg):
if msg.is_dead:
return
stream = self.stream_by_id(msg.src_id)
if stream.remote_id != msg.src_id or stream.detached:
LOG.warning('bad DETACHING received on %r: %r', stream, msg)
return
LOG.debug('%r: marking as detached', stream)
stream.detached = True
msg.reply(None)
def allocate_block(self):
self.lock.acquire()
try:
id_ = self.next_id
self.next_id += self.BLOCK_SIZE
end_id = id_ + self.BLOCK_SIZE
LOG.debug('%r: allocating [%d..%d)', self, id_, end_id)
return id_, end_id
finally:
self.lock.release()
def _send_load_module(self, stream, fullname):
if fullname not in stream.sent_modules:
LOG.debug('_send_load_module(%r, %r)', stream, fullname)
self._router._async_route(
mitogen.core.Message.pickled(
self._build_tuple(fullname),
dst_id=stream.remote_id,
handle=mitogen.core.LOAD_MODULE,
))
stream.sent_modules.add(fullname)
def connect(self):
LOG.debug('%r.connect()', self)
self.pid, fd = self.start_child()
self.name = '%s.%s' % (self.name_prefix, self.pid)
self.receive_side = mitogen.core.Side(self, fd)
self.transmit_side = mitogen.core.Side(self, os.dup(fd))
LOG.debug('%r.connect(): child process stdin/stdout=%r',
self, self.receive_side.fd)
self._connect_bootstrap()
def fetch(self, path, sender, msg):
"""
Start a transfer for a registered path.
:param str path:
File path.
:param mitogen.core.Sender sender:
Sender to receive file data.
:returns:
Dict containing the file metadata:
* ``size``: File size in bytes.
* ``mode``: Integer file mode.
* ``owner``: Owner account name on host machine.
* ``group``: Owner group name on host machine.
* ``mtime``: Floating point modification time.
* ``ctime``: Floating point change time.
:raises Error:
Unregistered path, or Sender did not match requestee context.
"""
if path not in self._paths and not self._prefix_is_authorized(path):
msg.reply(mitogen.core.CallError(
Error(self.unregistered_msg % (path,))
))
return
if msg.src_id != sender.context.context_id:
msg.reply(mitogen.core.CallError(
Error(self.context_mismatch_msg)
))
return
LOG.debug('Serving %r', path)
# Response must arrive first so requestee can begin receive loop,
# otherwise first ack won't arrive until all pending chunks were
# delivered. In that case max BDP would always be 128KiB, aka. max
# ~10Mbit/sec over a 100ms link.
try:
fp = open(path, 'rb', self.IO_SIZE)
msg.reply(self._generate_stat(path))
except IOError:
msg.reply(mitogen.core.CallError(
sys.exc_info()[1]
))
return
stream = self.router.stream_by_id(sender.context.context_id)
state = self._state_by_stream.setdefault(stream, FileStreamState())
state.lock.acquire()
try:
state.jobs.append((sender, fp))
self._schedule_pending_unlocked(state)
finally:
state.lock.release()
def register_prefix(self, path):
"""
Authorize a path and any subpaths for access by children. Repeat calls
with the same path has no effect.
:param str path:
File path.
"""
if path not in self._prefixes:
LOG.debug('%r: registering prefix %r', self, path)
self._prefixes.add(path)
def register(self, path):
"""
Authorize a path for access by children. Repeat calls with the same
path has no effect.
:param str path:
File path.
"""
if path not in self._paths:
LOG.debug('%r: registering %r', self, path)
self._paths.add(path)
def _send_module_and_related(self, stream, fullname):
tup = self.importer._cache[fullname]
for related in tup[4]:
rtup = self.importer._cache.get(related)
if rtup:
self._send_one_module(stream, rtup)
else:
LOG.debug('%r._send_module_and_related(%r): absent: %r',
self, fullname, related)
self._send_one_module(stream, tup)
def on_allocate_id(self, msg):
if msg == mitogen.core._DEAD:
return
id_, last_id = self.allocate_block()
requestee = self.router.context_by_id(msg.src_id)
allocated = self.router.context_by_id(id_, msg.src_id)
LOG.debug('%r: allocating [%r..%r) to %r', self, allocated, requestee,
msg.src_id)
msg.reply((id_, last_id))
def _on_get_module(self, msg):
if msg.is_dead:
return
LOG.debug('%r._on_get_module(%r)', self, msg.data)
stream = self._router.stream_by_id(msg.src_id)
fullname = msg.data.decode()
if fullname in stream.sent_modules:
LOG.warning('_on_get_module(): dup request for %r from %r',
fullname, stream)
self._send_module_and_related(stream, fullname)
def forward(self, path, context):
LOG.debug('%r.forward(%r, %r)', self, path, context)
func = lambda: self._forward(context, path)
self._lock.acquire()
try:
if path in self._cache:
func()
else:
LOG.debug('%r: %r not cached yet, queueing', self, path)
self._waiters.setdefault(path, []).append(func)
finally:
self._lock.release()
def _on_stream_disconnect(self, stream):
"""
Respond to disconnection of a local stream by
"""
LOG.debug('%r is gone; propagating DEL_ROUTE for %r',
stream, stream.routes)
for target_id in stream.routes:
self.router.del_route(target_id)
self.propagate(mitogen.core.DEL_ROUTE, target_id)
context = self.router.context_by_id(target_id, create=False)
if context:
mitogen.core.fire(context, 'disconnect')
def _on_cache_callback(self, msg, fullname):
LOG.debug('%r._on_get_module(): sending %r', self, fullname)
tup = self.importer._cache[fullname]
if tup is not None:
for related in tup[4]:
rtup = self.importer._cache.get(related)
if not rtup:
LOG.debug('%r._on_get_module(): skipping absent %r', self,
related)
continue
self._send_one_module(msg, rtup)
self._send_one_module(msg, tup)
def __init__(self, router, path=None, backlog=100):
self._router = router
self.path = path or make_socket_path()
self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
if os.path.exists(self.path) and is_path_dead(self.path):
LOG.debug('%r: deleting stale %r', self, self.path)
os.unlink(self.path)
self._sock.bind(self.path)
os.chmod(self.path, int('0600', 8))
self._sock.listen(backlog)
self.receive_side = mitogen.core.Side(self, self._sock.fileno())
router.broker.start_receive(self)
def _fakessh_main(dest_context_id, econtext):
hostname, opts, args = parse_args()
if not hostname:
die('Missing hostname')
subsystem = False
for opt, optarg in opts:
if opt == '-s':
subsystem = True
else:
LOG.debug('Warning option %s %s is ignored.', opt, optarg)
LOG.debug('hostname: %r', hostname)
LOG.debug('opts: %r', opts)
LOG.debug('args: %r', args)
if subsystem:
die('-s <subsystem> is not yet supported')
if not args:
die('fakessh: login mode not supported and no command specified')
dest = mitogen.parent.Context(econtext.router, dest_context_id)
# Even though SSH receives an argument vector, it still cats the vector
# together before sending to the server, the server just uses /bin/sh -c to
# run the command. We must remain puke-for-puke compatible.
control_handle, stdin_handle = dest.call(_start_slave,
mitogen.context_id, ' '.join(args))
LOG.debug('_fakessh_main: received control_handle=%r, stdin_handle=%r',
control_handle, stdin_handle)
process = Process(econtext.router, 1, 0)
process.start_master(
stdin=mitogen.core.Sender(dest, stdin_handle),
control=mitogen.core.Sender(dest, control_handle),
)
process.wait()
process.control.put(('exit', None))
def _on_pump_disconnect(self):
LOG.debug('%r._on_pump_disconnect()', self)
mitogen.core.fire(self, 'disconnect')
self.stdin.close()
self.wake_event.set()
def _on_proc_exit(self, status):
LOG.debug('%r._on_proc_exit(%r)', self, status)
self.control.put(('exit', status))
def _on_exit(self, msg, arg):
LOG.debug('on_exit: proc = %r', self.proc)
if self.proc:
self.proc.terminate()
else:
self.router.broker.shutdown()
