websockify.py

#!/usr/bin/env python


import sys
import asyncio
import websockets


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, **kwargs):
        return websockets.serve(self._server_run, bind_address, bind_port,
                                loop=self.loop, **kwargs)

    async def _server_run(self, websocket, path):
        print("WS connected.")
        try:
            (reader, writer) = await asyncio.open_connection(self.host, self.port,
                                                             loop=self.loop, **self.kwargs)
            self.loop.create_task(self._sock_run(websocket, reader))
            try:
                while True:
                    data = await websocket.recv()
                    if not data:
                        break
                    writer.write(data)
            except:
                pass
            try:
                writer.close()
            except:
                pass
        except Exception as err:
            await websocket.close()
            print(err)
        print("WS disconnected.")

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


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


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

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