344 lines
9.1 KiB
Python
344 lines
9.1 KiB
Python
import select
|
|
import socket
|
|
import struct
|
|
import sys
|
|
import types
|
|
|
|
VERBOSE = None
|
|
|
|
class SocketProtocol:
|
|
"""A simple protocol for sending strings across a socket"""
|
|
BUF_SIZE = 8192
|
|
|
|
def __init__(self, sock):
|
|
self.sock = sock
|
|
self._buffer = ''
|
|
self._closed = 0
|
|
|
|
def close(self):
|
|
self._closed = 1
|
|
self.sock.close()
|
|
|
|
def send(self, buf):
|
|
"""Encode buf and write it on the socket"""
|
|
if VERBOSE:
|
|
VERBOSE.write('send %d:%s\n' % (len(buf), `buf`))
|
|
self.sock.send('%d:%s' % (len(buf), buf))
|
|
|
|
def receive(self, timeout=0):
|
|
"""Get next complete string from socket or return None
|
|
|
|
Raise EOFError on EOF
|
|
"""
|
|
buf = self._read_from_buffer()
|
|
if buf is not None:
|
|
return buf
|
|
recvbuf = self._read_from_socket(timeout)
|
|
if recvbuf is None:
|
|
return None
|
|
if recvbuf == '' and self._buffer == '':
|
|
raise EOFError
|
|
if VERBOSE:
|
|
VERBOSE.write('recv %s\n' % `recvbuf`)
|
|
self._buffer = self._buffer + recvbuf
|
|
r = self._read_from_buffer()
|
|
return r
|
|
|
|
def _read_from_socket(self, timeout):
|
|
"""Does not block"""
|
|
if self._closed:
|
|
return ''
|
|
if timeout is not None:
|
|
r, w, x = select.select([self.sock], [], [], timeout)
|
|
if timeout is None or r:
|
|
return self.sock.recv(self.BUF_SIZE)
|
|
else:
|
|
return None
|
|
|
|
def _read_from_buffer(self):
|
|
buf = self._buffer
|
|
i = buf.find(':')
|
|
if i == -1:
|
|
return None
|
|
buflen = int(buf[:i])
|
|
enclen = i + 1 + buflen
|
|
if len(buf) >= enclen:
|
|
s = buf[i+1:enclen]
|
|
self._buffer = buf[enclen:]
|
|
return s
|
|
else:
|
|
self._buffer = buf
|
|
return None
|
|
|
|
# helpers for registerHandler method below
|
|
|
|
def get_methods(obj):
|
|
methods = []
|
|
for name in dir(obj):
|
|
attr = getattr(obj, name)
|
|
if callable(attr):
|
|
methods.append(name)
|
|
if type(obj) == types.InstanceType:
|
|
methods = methods + get_methods(obj.__class__)
|
|
if type(obj) == types.ClassType:
|
|
for super in obj.__bases__:
|
|
methods = methods + get_methods(super)
|
|
return methods
|
|
|
|
class CommandProtocol:
|
|
def __init__(self, sockp):
|
|
self.sockp = sockp
|
|
self.seqno = 0
|
|
self.handlers = {}
|
|
|
|
def close(self):
|
|
self.sockp.close()
|
|
self.handlers.clear()
|
|
|
|
def registerHandler(self, handler):
|
|
"""A Handler is an object with handle_XXX methods"""
|
|
for methname in get_methods(handler):
|
|
if methname[:7] == "handle_":
|
|
name = methname[7:]
|
|
self.handlers[name] = getattr(handler, methname)
|
|
|
|
def send(self, cmd, arg='', seqno=None):
|
|
if arg:
|
|
msg = "%s %s" % (cmd, arg)
|
|
else:
|
|
msg = cmd
|
|
if seqno is None:
|
|
seqno = self.get_seqno()
|
|
msgbuf = self.encode_seqno(seqno) + msg
|
|
self.sockp.send(msgbuf)
|
|
if cmd == "reply":
|
|
return
|
|
reply = self.sockp.receive(timeout=None)
|
|
r_cmd, r_arg, r_seqno = self._decode_msg(reply)
|
|
assert r_seqno == seqno and r_cmd == "reply", "bad reply"
|
|
return r_arg
|
|
|
|
def _decode_msg(self, msg):
|
|
seqno = self.decode_seqno(msg[:self.SEQNO_ENC_LEN])
|
|
msg = msg[self.SEQNO_ENC_LEN:]
|
|
parts = msg.split(" ", 2)
|
|
if len(parts) == 1:
|
|
cmd = msg
|
|
arg = ''
|
|
else:
|
|
cmd = parts[0]
|
|
arg = parts[1]
|
|
return cmd, arg, seqno
|
|
|
|
def dispatch(self):
|
|
msg = self.sockp.receive()
|
|
if msg is None:
|
|
return
|
|
cmd, arg, seqno = self._decode_msg(msg)
|
|
self._current_reply = seqno
|
|
h = self.handlers.get(cmd, self.default_handler)
|
|
try:
|
|
r = h(arg)
|
|
except TypeError, msg:
|
|
raise TypeError, "handle_%s: %s" % (cmd, msg)
|
|
if self._current_reply is None:
|
|
if r is not None:
|
|
sys.stderr.write("ignoring %s return value type %s\n" % \
|
|
(cmd, type(r).__name__))
|
|
return
|
|
if r is None:
|
|
r = ''
|
|
if type(r) != types.StringType:
|
|
raise ValueError, "invalid return type for %s" % cmd
|
|
self.send("reply", r, seqno=seqno)
|
|
|
|
def reply(self, arg=''):
|
|
"""Send a reply immediately
|
|
|
|
otherwise reply will be sent when handler returns
|
|
"""
|
|
self.send("reply", arg, self._current_reply)
|
|
self._current_reply = None
|
|
|
|
def default_handler(self, arg):
|
|
sys.stderr.write("WARNING: unhandled message %s\n" % arg)
|
|
return ''
|
|
|
|
SEQNO_ENC_LEN = 4
|
|
|
|
def get_seqno(self):
|
|
seqno = self.seqno
|
|
self.seqno = seqno + 1
|
|
return seqno
|
|
|
|
def encode_seqno(self, seqno):
|
|
return struct.pack("I", seqno)
|
|
|
|
def decode_seqno(self, buf):
|
|
return struct.unpack("I", buf)[0]
|
|
|
|
|
|
class StdioRedirector:
|
|
"""Redirect sys.std{in,out,err} to a set of file-like objects"""
|
|
|
|
def __init__(self, stdin, stdout, stderr):
|
|
self.stdin = stdin
|
|
self.stdout = stdout
|
|
self.stderr = stderr
|
|
|
|
def redirect(self):
|
|
self.save()
|
|
sys.stdin = self.stdin
|
|
sys.stdout = self.stdout
|
|
sys.stderr = self.stderr
|
|
|
|
def save(self):
|
|
self._stdin = sys.stdin
|
|
self._stdout = sys.stdout
|
|
self._stderr = sys.stderr
|
|
|
|
def restore(self):
|
|
sys.stdin = self._stdin
|
|
sys.stdout = self._stdout
|
|
sys.stderr = self._stderr
|
|
|
|
class IOWrapper:
|
|
"""Send output from a file-like object across a SocketProtocol
|
|
|
|
XXX Should this be more tightly integrated with the CommandProtocol?
|
|
"""
|
|
|
|
def __init__(self, name, cmdp):
|
|
self.name = name
|
|
self.cmdp = cmdp
|
|
self.buffer = []
|
|
|
|
class InputWrapper(IOWrapper):
|
|
def write(self, buf):
|
|
# XXX what should this do on Windows?
|
|
raise IOError, (9, '[Errno 9] Bad file descriptor')
|
|
|
|
def read(self, arg=None):
|
|
if arg is not None:
|
|
if arg <= 0:
|
|
return ''
|
|
else:
|
|
arg = 0
|
|
return self.cmdp.send(self.name, "read,%s" % arg)
|
|
|
|
def readline(self):
|
|
return self.cmdp.send(self.name, "readline")
|
|
|
|
class OutputWrapper(IOWrapper):
|
|
def write(self, buf):
|
|
self.cmdp.send(self.name, buf)
|
|
|
|
def read(self, arg=None):
|
|
return ''
|
|
|
|
class RemoteInterp:
|
|
def __init__(self, sock):
|
|
self._sock = SocketProtocol(sock)
|
|
self._cmd = CommandProtocol(self._sock)
|
|
self._cmd.registerHandler(self)
|
|
|
|
def run(self):
|
|
try:
|
|
while 1:
|
|
self._cmd.dispatch()
|
|
except EOFError:
|
|
pass
|
|
|
|
def handle_execfile(self, arg):
|
|
self._cmd.reply()
|
|
io = StdioRedirector(InputWrapper("stdin", self._cmd),
|
|
OutputWrapper("stdout", self._cmd),
|
|
OutputWrapper("stderr", self._cmd))
|
|
io.redirect()
|
|
execfile(arg, {'__name__':'__main__'})
|
|
io.restore()
|
|
self._cmd.send("terminated")
|
|
|
|
def handle_quit(self, arg):
|
|
self._cmd.reply()
|
|
self._cmd.close()
|
|
|
|
def startRemoteInterp(id):
|
|
import os
|
|
# UNIX domain sockets are simpler for starters
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
sock.bind("/var/tmp/ri.%s" % id)
|
|
try:
|
|
sock.listen(1)
|
|
cli, addr = sock.accept()
|
|
rinterp = RemoteInterp(cli)
|
|
rinterp.run()
|
|
finally:
|
|
os.unlink("/var/tmp/ri.%s" % id)
|
|
|
|
class RIClient:
|
|
"""Client of the remote interpreter"""
|
|
def __init__(self, sock):
|
|
self._sock = SocketProtocol(sock)
|
|
self._cmd = CommandProtocol(self._sock)
|
|
self._cmd.registerHandler(self)
|
|
|
|
def execfile(self, file):
|
|
self._cmd.send("execfile", file)
|
|
|
|
def run(self):
|
|
try:
|
|
while 1:
|
|
self._cmd.dispatch()
|
|
except EOFError:
|
|
pass
|
|
|
|
def handle_stdout(self, buf):
|
|
sys.stdout.write(buf)
|
|
## sys.stdout.flush()
|
|
|
|
def handle_stderr(self, buf):
|
|
sys.stderr.write(buf)
|
|
|
|
def handle_stdin(self, arg):
|
|
if arg == "readline":
|
|
return sys.stdin.readline()
|
|
i = arg.find(",") + 1
|
|
bytes = int(arg[i:])
|
|
if bytes == 0:
|
|
return sys.stdin.read()
|
|
else:
|
|
return sys.stdin.read(bytes)
|
|
|
|
def handle_terminated(self, arg):
|
|
self._cmd.reply()
|
|
self._cmd.send("quit")
|
|
self._cmd.close()
|
|
|
|
def riExec(id, file):
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
sock.connect("/var/tmp/ri.%s" % id)
|
|
cli = RIClient(sock)
|
|
cli.execfile(file)
|
|
cli.run()
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
import getopt
|
|
|
|
SERVER = 1
|
|
opts, args = getopt.getopt(sys.argv[1:], 'cv')
|
|
for o, v in opts:
|
|
if o == '-c':
|
|
SERVER = 0
|
|
elif o == '-v':
|
|
VERBOSE = sys.stderr
|
|
id = args[0]
|
|
|
|
if SERVER:
|
|
startRemoteInterp(id)
|
|
else:
|
|
file = args[1]
|
|
riExec(id, file)
|
|
|