websockify2.py

#!/usr/bin/env python


import sys
import asyncio
from aiohttp import web


class Server:
    def __init__(self, host, port, loop=None, **kwargs):
        super().__init__()

        self.host = host
        self.port = port
        self.loop = loop
        self.kwargs = kwargs

        if not self.loop:
            self.loop = asyncio.get_event_loop()

    def listen(self, bind_address, bind_port, bind_path, **kwargs):
        app = web.Application()
        app.router.add_get(bind_path, self._ws_handler)
        web.run_app(app, host=bind_address, port=int(bind_port))

    async def _ws_handler(self, request):
        print("WS connected.")
        try:
            ws = web.WebSocketResponse()
            await ws.prepare(request)

            (reader, writer) = await asyncio.open_connection(self.host, self.port,
                                                             loop=self.loop, **self.kwargs)
            self.loop.create_task(self._sock_run(ws, reader))
            try:
                while True:
                    data = await ws.receive_bytes()
                    if not data:
                        break
                    writer.write(data)
            except:
                pass
            try:
                writer.close()
            except:
                pass
        except Exception as err:
            await ws.close()
            print(err)
        print("WS disconnected.")
        return ws

    async def _sock_run(self, ws, reader):
        print("Socket connected.")
        try:
            while True:
                data = await reader.read(65536)
                if not data:
                    break
                await ws.send_bytes(data)
        except:
            pass
        await ws.close()
        print("Socket disconnected")


def _usage():
    print("Usage: {} local_host local_port ws_path remote_host remote_port".format(sys.argv[0]))


if __name__ == '__main__':
    if len(sys.argv) != 6:
        _usage()
        exit(1)

    loop = asyncio.get_event_loop()
    server = Server(sys.argv[4], sys.argv[5])
    server.listen(sys.argv[1], sys.argv[2], sys.argv[3])
    loop.run_forever()
    loop.close()