import socket
import sys
import importlib
import io
import traceback
class INST:
UNKNOWN = 0x00
PRINT = 0x01
LOAD = 0x02
EXCEPTION = 0x03
INITIALIZE = 0x04
EXIT = 0x05
EXECUTE = 0x06
class MessageStream:
def __init__(self, socket):
self.socket = socket
self._read_buf = bytearray()
def recv_msg(self):
self._read_buf.clear()
self._read_buf.extend(self.socket.recv(3))
inst = int.from_bytes(self._read_buf[:1], 'big')
data_len = int.from_bytes(self._read_buf[1:3], 'big')
self._read_buf.extend(self.socket.recv(data_len))
return (inst, self._read_buf[3:].decode('utf-8'))
def send_msg(self, inst, data=''):
data_bytes = data.encode()
data_len = len(data_bytes)
raw_bytes = inst.to_bytes(1, 'big') + data_len.to_bytes(2, 'big') + data_bytes
self.socket.send(raw_bytes)
def close(self):
self.socket.close()
server_socket = socket.socket()
server_socket.bind(('127.0.0.1', __PORT__))
server_socket.listen(1)
(client_socket, client_address) = server_socket.accept()
already_loaded = False
ctx = {'importlib': importlib}
client_stream = MessageStream(client_socket)
while True:
try:
inst, data = client_stream.recv_msg()
except ConnectionResetError: break
if inst == INST.EXIT: client_stream.send_msg(INST.EXIT)
break
elif inst == INST.LOAD:
sys.stdout = io.StringIO()
res = ''
exc = ''
resp_inst = INST.PRINT
buf = []
try:
if already_loaded:
res = str(exec('importlib.reload(__MODULE__)', ctx))
else:
res = str(exec('import __MODULE__', ctx))
already_loaded = True
except SystemExit:
client_stream.send_msg(INST.EXCEPTION, 'SystemExit')
continue
except Exception as e:
try:
excs = traceback.format_exception(e)
except:
excs = traceback.format_exception_only(e.__class__, e)
exc = ''.join(excs).rstrip()
traceback.clear_frames(e.__traceback__)
resp_inst = INST.INITIALIZE
out = sys.stdout.getvalue()[:-1]
if out and exc or res:
out += '\n'
res = out + exc + res
buf.append(res)
client_stream.send_msg(resp_inst, ''.join(buf))
elif inst == INST.EXECUTE:
sys.stdout = io.StringIO()
res = ''
exc = ''
resp_inst = INST.PRINT
buf = []
try:
res = str(exec(data, ctx))
except SystemExit:
client_stream.send_msg(INST.EXCEPTION, 'SystemExit')
continue
except Exception as e:
try:
excs = traceback.format_exception(e)
except:
excs = traceback.format_exception_only(e.__class__, e)
exc = ''.join(excs).rstrip()
traceback.clear_frames(e.__traceback__)
resp_inst = INST.INITIALIZE
out = sys.stdout.getvalue()[:-1]
if out and exc or res:
out += '\n'
res = out + exc + res
buf.append(res)
client_stream.send_msg(resp_inst, ''.join(buf))
else:
client_stream.send_msg(INST.UNKNOWN)
client_stream.close()
server_socket.close()