japronto_ws.py

import asyncio
import base64
import hashlib

from japronto.app import Application


class NaiveWsBaseProtocol(asyncio.Protocol):
    @staticmethod
    def accept(key):
        GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
        sha1 = hashlib.sha1((key + GUID).encode()).digest()
        return base64.b64encode(sha1).decode()

    def __init__(self, transport, key):
        self.eof_feature = asyncio.Future()
        self.transport = transport
        self.transport.set_protocol(self)
        self.key = key

    async def handshake(self):
        response = ['HTTP/1.1 101 Switching Protocols']
        response.extend('{}: {}'.format(k, v) for k, v in [
            ('Upgrade', 'WebSocket'),
            ('Connection', 'Upgrade'),
            ('Sec-WebSocket-Accept', NaiveWsBaseProtocol.accept(self.key))
        ])
        response.append('\r\n')
        response = '\r\n'.join(response).encode()

        self.transport.write(response)

    def data_received(self, data):
        print(data)

    def eof_received(self):
        print('eof received')
        self.eof_feature.set_result(None)


async def ws_handler(request):
    print('Common HTTP\n{}'.format(
        '\n'.join('\t{}: {}'.format(k, v) for k, v in request.headers.items())
    ))
    ws = NaiveWsBaseProtocol(request.transport, request.headers.get('Sec-Websocket-Key'))
    await ws.handshake()
    await ws.eof_feature


app = Application()
app.router.add_route('/ws', ws_handler)
app.run(port=10000)