54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
from .engines import engines
|
|
from .helpers import ProxyProtocolMixin, ZlibPacketHandler
|
|
|
|
|
|
class EchoPacketHandler(ProxyProtocolMixin, ZlibPacketHandler):
|
|
def __init__(self, server, socket):
|
|
super(EchoPacketHandler, self).__init__(server, socket)
|
|
self._gotdata = False
|
|
self.server.schedule(5, self._kill_if_no_data)
|
|
|
|
def _kill_if_no_data(self):
|
|
if not self._gotdata:
|
|
print('Inactive client:', self._socket.getpeername())
|
|
self.close()
|
|
|
|
def packet(self, data):
|
|
self._gotdata = True
|
|
print('Data from %s: %r' % (self._socket.getpeername(), data[:30] if len(data) > 30 else data))
|
|
self.send(data)
|
|
|
|
def on_close(self):
|
|
self._gotdata = True
|
|
print('Closed client:', self._socket.getpeername())
|
|
|
|
|
|
def main():
|
|
import argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('-l', '--host', action='append')
|
|
parser.add_argument('-p', '--port', type=int, action='append')
|
|
parser.add_argument('-e', '--engine', default='select', choices=sorted(engines.keys()))
|
|
try:
|
|
import netaddr
|
|
except ImportError:
|
|
netaddr = None
|
|
else:
|
|
parser.add_argument('-P', '--proxy', action='append')
|
|
args = parser.parse_args()
|
|
|
|
class TestServer(engines[args.engine]):
|
|
def _accept(self, sock):
|
|
client = super(TestServer, self)._accept(sock)
|
|
print('New connection:', client.socket.getpeername())
|
|
return client
|
|
|
|
handler = EchoPacketHandler
|
|
if netaddr is not None and args.proxy:
|
|
handler = handler.with_proxy_set(args.proxy)
|
|
server = TestServer(list(zip(args.host, args.port)), handler)
|
|
server.serve_forever()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|