Fork of `rfoo.utils.rconsole` to run with other RPC implementations.

remconsole.py

"""
    remconsole.py

    A Python console you can embed in a program and attach to remotely.

    To spawn a Python console in a script do the following in any scope
    of any module:

        import remconsole
        remconsole.spawn_server()

    This will start a listener for connections in a new thread. You may
    specify a port to listen on.

    To attach to the console from another shell execute this script or
    simply invoke interact().

    SECURITY NOTE:
    The listener started with spawn_server() will accept any local
    connection and may therefore be insecure to use in shared hosting
    or similar environments!
"""

import rlcompleter
import logging
import pprint
import code
import sys

logging.getLogger().setLevel(logging.INFO)

try:
    import thread
except ImportError:
    import _thread as thread


HANDLER_CLASS = object


try:
    import rpyc
    RPC_HANDLER = 'rpyc'
except ImportError:
    try:
        import Pyro.core
        RPC_HANDLER = 'Pyro'
    except ImportError:
        RPC_HANDLER = 'xmlrpc'

logging.info('RPC Handler is: %s.' % RPC_HANDLER)


if RPC_HANDLER == 'rpyc':
    from rpyc.utils.server import ThreadedServer

    def start_server(port, handler):
        class Service(rpyc.Service):

            def exposed_complete(self, *args, **kwargs):
                return handler.complete(*args, **kwargs)

            def exposed_runsource(self, *args, **kwargs):
                return handler.runsource(*args, **kwargs)

        t = ThreadedServer(Service, port=port)
        t.start()

    def get_proxy(port):
        conn = rpyc.connect('localhost', port)
        rpyc.__remconsole_conn = conn  # Workaround to prevent GC
        return conn.root

elif RPC_HANDLER == 'Pyro':
    HANDLER_CLASS = Pyro.core.ObjBase

    def start_server(port, handler):
        daemon = Pyro.core.Daemon(port=port)
        daemon.connectPersistent(handler, 'console')
        daemon.requestLoop()

    def _auto_rebind(proxy, func):
        def wrapper(*args, **kwargs):
            while True:
                try:
                    return func(*args, **kwargs)
                except Pyro.errors.ConnectionClosedError:
                    logging.warning('Connection lost. REBINDING...\n'
                                    'This could hang forever. Quit with '
                                    'Ctrl+C in this case.')
                    proxy.adapter.rebindURI()
        return wrapper

    def get_proxy(port):
        uri = 'PYROLOC://localhost:%s/console' % port
        proxy = Pyro.core.getProxyForURI(uri)

        func = _auto_rebind(proxy, proxy.adapter.remoteInvocation)
        proxy.adapter.remoteInvocation = func
        return proxy

else:
    import xmlrpclib
    from SimpleXMLRPCServer import SimpleXMLRPCServer

    def start_server(port, handler):
        server = SimpleXMLRPCServer(('localhost', port), allow_none=True)
        server.register_function(handler.complete)
        server.register_function(handler.runsource)
        server.serve_forever()

    def get_proxy(port):
        return xmlrpclib.ServerProxy('http://localhost:%s/' % port)


PORT = 54321


class BufferedInterpreter(code.InteractiveInterpreter):
    """Variation of code.InteractiveInterpreter that outputs to buffer."""

    def __init__(self, *args, **kwargs):
        code.InteractiveInterpreter.__init__(self, *args, **kwargs)
        self.buffout = ''

    def write(self, data):
        self.buffout += data


class ConsoleHandler(HANDLER_CLASS):
    """An handler that remotes a Python interpreter."""

    def __init__(self, namespace, *args, **kwargs):
        super(ConsoleHandler, self).__init__(*args, **kwargs)

        self._namespace = namespace
        self._interpreter = BufferedInterpreter(self._namespace)
        self._completer = rlcompleter.Completer(self._namespace)

    def complete(self, phrase, state):
        """Auto complete for remote console."""
        logging.debug('Enter, phrase=%r, state=%d.', phrase, state)
        return self._completer.complete(phrase, state)

    def runsource(self, source, filename="<input>"):
        """Variation of InteractiveConsole which returns expression
        result as second element of returned tuple.
        """

        logging.debug('Enter, source=%r.', source)

        # Inject a global variable to capture expression result.
        self._namespace['_rcon_result_'] = None

        try:
            # In case of an expression, capture result.
            compile(source, '<input>', 'eval')
            source = '_rcon_result_ = ' + source
            logging.debug('source is an expression.')
        except SyntaxError:
            pass

        more = self._interpreter.runsource(source, filename)
        result = self._namespace.pop('_rcon_result_')

        if more is True:
            logging.debug('source is incomplete.')
            return True, ''

        output = self._interpreter.buffout
        self._interpreter.buffout = ''

        if result is not None:
            result = pprint.pformat(result)
            output += result + '\n'

        return False, output


class ProxyConsole(code.InteractiveConsole):
    """Proxy interactive console to remote interpreter."""

    def __init__(self, proxy):
        code.InteractiveConsole.__init__(self)
        self.proxy = proxy

    def interact(self, banner=None):
        logging.info('Enter.')

        return code.InteractiveConsole.interact(self, banner)

    def complete(self, phrase, state):
        """Auto complete support for interactive console."""

        logging.debug('Enter, phrase=%r, state=%d.', phrase, state)

        # Allow tab key to simply insert spaces when proper.
        if phrase == '':
            if state == 0:
                return '    '
            return None

        return self.proxy.complete(phrase, state)

    def runsource(self, source, filename="<input>", symbol="single"):
        logging.debug('Enter, source=%r.', source)

        more, output = self.proxy.runsource(source, filename)
        if output:
            self.write(output)

        return more


def spawn_server(namespace=None, port=PORT):
    """Start console server in a new thread.
    Should be called from global scope only!
    May be insecure on shared machines.
    """

    logging.info('Enter, port=%d.', port)

    if namespace is None:
        namespace = sys._getframe(1).f_globals.copy()
        namespace.update(sys._getframe(1).f_locals)

    handler = ConsoleHandler(namespace)
    thread.start_new_thread(start_server, (port, handler))


def interact(banner=None, readfunc=None, port=PORT):
    """Start console and connect to remote console server."""

    logging.info('Enter, port=%d.', port)
    proxy = get_proxy(port)
    console = ProxyConsole(proxy)

    if readfunc is not None:
        console.raw_input = readfunc
    else:
        try:
            import readline
            readline.set_completer(console.complete)
            readline.parse_and_bind('tab: complete')
        except ImportError:
            pass

    console.interact(banner)


if __name__ == '__main__':
    interact()