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 _accept_client(self, sock): sock.setblocking(True) try: pid, = struct.unpack('>L', sock.recv(4)) except socket.error: LOG.error('%r: failed to read remote identity: %s', self, sys.exc_info()[1]) return context_id = self._router.id_allocator.allocate() context = mitogen.parent.Context(self._router, context_id) stream = mitogen.core.Stream(self._router, context_id) stream.name = u'unix_client.%d' % (pid,) stream.auth_id = mitogen.context_id stream.is_privileged = True try: sock.send(struct.pack('>LLL', context_id, mitogen.context_id, os.getpid())) except socket.error: LOG.error('%r: failed to assign identity to PID %d: %s', self, pid, sys.exc_info()[1]) return stream.accept(sock.fileno(), sock.fileno()) self._router.register(context, stream)
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_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 __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 on_shutdown(self): """ Respond to shutdown by sending close() to every target, allowing their receive loop to exit and clean up gracefully. """ LOG.debug('%r.on_shutdown()', self) for stream, state in self._state_by_stream.items(): state.lock.acquire() try: for sender, fp in reversed(state.jobs): sender.close() fp.close() state.jobs.pop() finally: state.lock.release()
def store_and_forward(self, path, data, context): LOG.debug('%r.store_and_forward(%r, %r, %r) %r', self, path, data, context, threading.currentThread().getName()) self._lock.acquire() try: self._cache[path] = data waiters = self._waiters.pop(path, []) finally: self._lock.release() if context.context_id != mitogen.context_id: self._forward(context, path) for callback in waiters: callback()
def neutralize_main(self, path, src): """ Given the source for the __main__ module, try to find where it begins conditional execution based on a "if __name__ == '__main__'" guard, and remove any code after that point. """ match = self.MAIN_RE.search(src) if match: return src[:match.start()] if b('mitogen.main(') in src: return src LOG.error(self.main_guard_msg, path) raise ImportError('refused')
def _on_service_call(self, recv, msg): self._validate(msg) service_name, method_name, kwargs = msg.unpickle() try: invoker = self.get_invoker(service_name, msg) return invoker.invoke(method_name, kwargs, msg) except mitogen.core.CallError: e = sys.exc_info()[1] LOG.warning('%r: call error: %s: %s', self, msg, e) msg.reply(e) except Exception: LOG.exception('%r: while invoking %r of %r', self, method_name, service_name) e = sys.exc_info()[1] msg.reply(mitogen.core.CallError(e))
def propagate_route(self, target, via): self.add_route(target.context_id, via.context_id) child = via parent = via.via while parent is not None: LOG.debug('Adding route to %r for %r via %r', parent, target, child) parent.send( mitogen.core.Message( data='%s\x00%s' % (target.context_id, child.context_id), handle=mitogen.core.ADD_ROUTE, )) child = parent parent = parent.via
def get(self, path): self._lock.acquire() try: if path in self._cache: return self._cache[path] waiters = self._waiters.setdefault(path, []) latch = mitogen.core.Latch() waiters.append(lambda: latch.put(None)) finally: self._lock.release() LOG.debug('%r.get(%r) waiting for uncached file to arrive', self, path) latch.get() LOG.debug('%r.get(%r) -> %r', self, path, self._cache[path]) return self._cache[path]
def _run(self): while True: tup = self._pop() if tup is None: return method_name, kwargs, msg = tup try: super(SerializedInvoker, self).invoke(method_name, kwargs, msg) except mitogen.core.CallError: e = sys.exc_info()[1] LOG.warning('%r: call error: %s: %s', self, msg, e) msg.reply(e) except Exception: LOG.exception('%r: while invoking %s()', self, method_name) msg.reply(mitogen.core.Message.dead())
def find(self, fullname): """ Find `fullname` using :func:`pkgutil.find_loader`. """ try: # Pre-'import spec' this returned None, in Python3.6 it raises # ImportError. loader = pkgutil.find_loader(fullname) except ImportError: e = sys.exc_info()[1] LOG.debug('%r._get_module_via_pkgutil(%r): %s', self, fullname, e) return None IOLOG.debug('%r._get_module_via_pkgutil(%r) -> %r', self, fullname, loader) if not loader: return try: path, is_special = _py_filename(loader.get_filename(fullname)) source = loader.get_source(fullname) is_pkg = loader.is_package(fullname) # workaround for special python modules that might only exist in memory if is_special and is_pkg and not source: source = '\n' except (AttributeError, ImportError): # - Per PEP-302, get_source() and is_package() are optional, # calling them may throw AttributeError. # - get_filename() may throw ImportError if pkgutil.find_loader() # picks a "parent" package's loader for some crap that's been # stuffed in sys.modules, for example in the case of urllib3: # "loader for urllib3.contrib.pyopenssl cannot handle # requests.packages.urllib3.contrib.pyopenssl" e = sys.exc_info()[1] LOG.debug('%r: loading %r using %r failed: %s', self, fullname, loader, e) return if path is None or source is None: return if isinstance(source, mitogen.core.UnicodeType): # get_source() returns "string" according to PEP-302, which was # reinterpreted for Python 3 to mean a Unicode string. source = source.encode('utf-8') return path, source, is_pkg
def _on_del_route(self, msg): if msg == mitogen.core._DEAD: return target_id = int(msg.data) registered_stream = self.router.stream_by_id(target_id) stream = self.router.stream_by_id(msg.auth_id) if registered_stream != stream: LOG.error('Received DEL_ROUTE for %d from %r, expected %r', target_id, stream, registered_stream) return LOG.debug('Deleting route to %d via %r', target_id, stream) stream.routes.discard(target_id) self.router.del_route(target_id) self.propagate(mitogen.core.DEL_ROUTE, target_id)
def build_stream(cls, router, path=None, backlog=100): if not path: path = make_socket_path() sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) if os.path.exists(path) and is_path_dead(path): LOG.debug('%r: deleting stale %r', cls.__name__, path) os.unlink(path) sock.bind(path) os.chmod(path, int('0600', 8)) sock.listen(backlog) stream = super(Listener, cls).build_stream(router, path) stream.accept(sock, sock) router.broker.start_receive(stream) return stream
def _on_forward_log(self, msg): if msg == mitogen.core._DEAD: return logger = self._cache.get(msg.src_id) if logger is None: context = self._router.context_by_id(msg.src_id) if context is None: LOG.error('FORWARD_LOG received from src_id %d', msg.src_id) return name = '%s.%s' % (RLOG.name, context.name) self._cache[msg.src_id] = logger = logging.getLogger(name) name, level_s, s = msg.data.split('\x00', 2) logger.log(int(level_s), '%s: %s', name, s)
def _send_load_module(self, stream, fullname): if fullname not in stream.sent_modules: LOG.debug('_send_load_module(%r, %r)', stream, fullname) tup = self._build_tuple(fullname) msg = mitogen.core.Message.pickled( tup, dst_id=stream.remote_id, handle=mitogen.core.LOAD_MODULE, ) self._router._async_route(msg) stream.sent_modules.add(fullname) if tup[2] is not None: self.good_load_module_count += 1 self.good_load_module_size += len(msg.data) else: self.bad_load_module_count += 1
def on_receive(self, broker): """ This handler is only called after the stream is registered with the IO loop, the descriptor is manually read/written by _connect_bootstrap() prior to that. """ buf = self.receive_side.read() if not buf: return self.on_disconnect(broker) self.buf += buf while '\n' in self.buf: lines = self.buf.split('\n') self.buf = lines[-1] for line in lines[:-1]: LOG.debug('%r: %r', self, line.rstrip())
def _send_load_module(self, stream, fullname): if fullname not in stream.protocol.sent_modules: tup = self._build_tuple(fullname) msg = mitogen.core.Message.pickled( tup, dst_id=stream.protocol.remote_id, handle=mitogen.core.LOAD_MODULE, ) LOG.debug('%s: sending %s (%.2f KiB) to %s', self, fullname, len(msg.data) / 1024.0, stream.name) self._router._async_route(msg) stream.protocol.sent_modules.add(fullname) if tup[2] is not None: self.good_load_module_count += 1 self.good_load_module_size += len(msg.data) else: self.bad_load_module_count += 1
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._metadata_by_path: raise Error(self.unregistered_msg) if msg.src_id != sender.context.context_id: raise Error(self.context_mismatch_msg) LOG.debug('Serving %r', path) try: fp = open(path, 'rb', self.IO_SIZE) except IOError: msg.reply(mitogen.core.CallError(sys.exc_info()[1])) return # 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. msg.reply(self._metadata_by_path[path]) 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 poll(self, timeout=None): changelist = self._changelist self._changelist = [] events, _ = mitogen.core.io_op(self._kqueue.control, changelist, 32, timeout) for event in events: fd = event.ident if event.flags & select.KQ_EV_ERROR: LOG.debug('ignoring stale event for fd %r: errno=%d: %s', fd, event.data, errno.errorcode.get(event.data)) elif event.filter == select.KQ_FILTER_READ and fd in self._rfds: # Events can still be read for an already-discarded fd. mitogen.core._vv and IOLOG.debug('%r: POLLIN: %r', self, fd) yield self._rfds[fd] elif event.filter == select.KQ_FILTER_WRITE and fd in self._wfds: mitogen.core._vv and IOLOG.debug('%r: POLLOUT: %r', self, fd) yield self._wfds[fd]
def _on_broker_exit(self): super(Router, self)._on_broker_exit() dct = self.get_stats() dct['self'] = self dct['get_module_ms'] = 1000 * dct['get_module_secs'] dct['good_load_module_size_kb'] = dct['good_load_module_size'] / 1024.0 dct['good_load_module_size_avg'] = ( (dct['good_load_module_size'] / (float(dct['good_load_module_count']) or 1.0)) / 1024.0) LOG.debug('%(self)r: stats: ' '%(get_module_count)d module requests in ' '%(get_module_ms)d ms, ' '%(good_load_module_count)d sent, ' '%(bad_load_module_count)d negative responses. ' 'Sent %(good_load_module_size_kb).01f kb total, ' '%(good_load_module_size_avg).01f kb avg.' % dct)
def __init__(self, router, path=None, backlog=30): self._router = router self.path = path or tempfile.mktemp(prefix='mitogen_unix_') 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, 0600) self._sock.listen(backlog) mitogen.core.set_nonblock(self._sock.fileno()) mitogen.core.set_cloexec(self._sock.fileno()) self.path = self._sock.getsockname() self.receive_side = mitogen.core.Side(self, self._sock.fileno()) router.broker.start_receive(self)
def _on_get_module(self, msg): if msg.is_dead: return LOG.debug('%r._on_get_module(%r)', self, msg.data) self.get_module_count += 1 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) t0 = time.time() try: self._send_module_and_related(stream, fullname) finally: self.get_module_secs += time.time() - t0
def acknowledge(self, size, msg): """ Acknowledge bytes received by a transfer target, scheduling new chunks to keep the window full. This should be called for every chunk received by the target. """ stream = self.router.stream_by_id(msg.src_id) state = self._state_by_stream[stream] state.lock.acquire() try: if state.unacked < size: LOG.error('%r.acknowledge(src_id %d): unacked=%d < size %d', self, msg.src_id, state.unacked, size) state.unacked -= min(state.unacked, size) self._schedule_pending_unlocked(state) finally: state.lock.release()
def _send_module_and_related(self, stream, fullname): if fullname in stream.protocol.sent_modules: return try: tup = self._build_tuple(fullname) for name in tup[4]: # related parent, _, _ = str_partition(name, '.') if parent != fullname and parent not in stream.protocol.sent_modules: # Parent hasn't been sent, so don't load submodule yet. continue self._send_load_module(stream, name) self._send_load_module(stream, fullname) except Exception: LOG.debug('While importing %r', fullname, exc_info=True) self._send_module_load_failed(stream, fullname)
def _get_module_via_pkgutil(self, fullname): """ Attempt to fetch source code via pkgutil. In an ideal world, this would be the only required implementation of get_module(). """ try: # Pre-'import spec' this returned None, in Python3.6 it raises # ImportError. loader = pkgutil.find_loader(fullname) except ImportError: e = sys.exc_info()[1] LOG.debug('%r._get_module_via_pkgutil(%r): %s', self, fullname, e) return None IOLOG.debug('%r._get_module_via_pkgutil(%r) -> %r', self, fullname, loader) if not loader: return try: path = self._py_filename(loader.get_filename(fullname)) source = loader.get_source(fullname) is_pkg = loader.is_package(fullname) except (AttributeError, ImportError): # - Per PEP-302, get_source() and is_package() are optional, # calling them may throw AttributeError. # - get_filename() may throw ImportError if pkgutil.find_loader() # picks a "parent" package's loader for some crap that's been # stuffed in sys.modules, for example in the case of urllib3: # "loader for urllib3.contrib.pyopenssl cannot handle # requests.packages.urllib3.contrib.pyopenssl" e = sys.exc_info()[1] LOG.debug('%r: loading %r using %r failed: %s', self, fullname, loader, e) return if path is None or source is None: return if isinstance(source, mitogen.core.UnicodeType): # get_source() returns "string" according to PEP-302, which was # reinterpreted for Python 3 to mean a Unicode string. source = source.encode('utf-8') return path, source, is_pkg
def _get_module_via_pkgutil(self, fullname): """Attempt to fetch source code via pkgutil. In an ideal world, this would be the only required implementation of get_module().""" loader = pkgutil.find_loader(fullname) LOG.debug('pkgutil._get_module_via_pkgutil(%r) -> %r', fullname, loader) if not loader: return try: path = self._py_filename(loader.get_filename(fullname)) source = loader.get_source(fullname) is_pkg = loader.is_package(fullname) except AttributeError: return if path is not None and source is not None: return path, source, is_pkg
def _invoke(self, method_name, kwargs, msg): method = getattr(self.service, method_name) if 'msg' in func_code(method).co_varnames: kwargs['msg'] = msg # TODO: hack no_reply = getattr(method, 'mitogen_service__no_reply', False) ret = None try: ret = method(**kwargs) if no_reply: return Service.NO_REPLY return ret except Exception: if no_reply: LOG.exception('While calling no-reply method %s.%s', type(self.service).__name__, func_name(method)) else: raise
def _on_receive_message(self, msg): method_name, kwargs = self._validate_message(msg) method = getattr(self, method_name) if 'msg' in method.func_code.co_varnames: kwargs['msg'] = msg # TODO: hack no_reply = getattr(method, 'mitogen_service__no_reply', False) ret = None try: ret = method(**kwargs) if no_reply: return self.NO_REPLY return ret except Exception: if no_reply: LOG.exception('While calling no-reply method %s.%s', type(self).__name__, method.func_name) else: raise
def create_child(*args): parentfp, childfp = socket.socketpair() pid = os.fork() if not pid: mitogen.core.set_block(childfp.fileno()) os.dup2(childfp.fileno(), 0) os.dup2(childfp.fileno(), 1) childfp.close() parentfp.close() os.execvp(args[0], args) childfp.close() # Decouple the socket from the lifetime of the Python socket object. fd = os.dup(parentfp.fileno()) parentfp.close() LOG.debug('create_child() child %d fd %d, parent %d, cmd: %s', pid, fd, os.getpid(), Argv(args)) return pid, fd
def get(self, path): """ Fetch a file from the cache. """ assert isinstance(path, mitogen.core.UnicodeType) self._lock.acquire() try: if path in self._cache: return self._cache[path] latch = mitogen.core.Latch() waiters = self._waiters.setdefault(path, []) waiters.append(lambda: latch.put(None)) finally: self._lock.release() LOG.debug('%r.get(%r) waiting for uncached file to arrive', self, path) latch.get() LOG.debug('%r.get(%r) -> %r', self, path, self._cache[path]) return self._cache[path]
def allocate_block(self): """ Allocate a block of IDs for use in a child context. This function is safe to call from any thread. :returns: Tuple of the form `(id, end_id)` where `id` is the first usable ID and `end_id` is the last usable ID. """ 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 _on_del_route(self, msg): if msg.is_dead: return target_id = int(msg.data) registered_stream = self.router.stream_by_id(target_id) stream = self.router.stream_by_id(msg.auth_id) if registered_stream != stream: LOG.error('Received DEL_ROUTE for %d from %r, expected %r', target_id, stream, registered_stream) return LOG.debug('Deleting route to %d via %r', target_id, stream) stream.routes.discard(target_id) 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_add_route(self, msg): if msg.is_dead: return target_id_s, _, target_name = msg.data.partition(':') target_id = int(target_id_s) self.router.context_by_id(target_id).name = target_name stream = self.router.stream_by_id(msg.auth_id) current = self.router.stream_by_id(target_id) if current and current.remote_id != mitogen.parent_id: LOG.error('Cannot add duplicate route to %r via %r, ' 'already have existing route via %r', target_id, stream, current) return LOG.debug('Adding route to %d via %r', target_id, stream) stream.routes.add(target_id) self.router.add_route(target_id, stream) self.propagate(mitogen.core.ADD_ROUTE, target_id, target_name)
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_exit(self, msg, arg): LOG.debug('on_exit: proc = %r', self.proc) if self.proc: self.proc.terminate() else: self.router.broker.shutdown()
def _on_proc_exit(self, status): LOG.debug('%r._on_proc_exit(%r)', self, status) self.control.put(('exit', status))
def _on_pump_disconnect(self): LOG.debug('%r._on_pump_disconnect()', self) mitogen.core.fire(self, 'disconnect') self.stdin.close() self.wake_event.set()