pico.chat/pico.py

import asyncio
import collections
import datetime
import http
import json
import functools
from pathlib import Path

import websockets
from websockets.headers import (
    build_www_authenticate_basic,
    parse_authorization_basic,
)


rooms = collections.defaultdict(set)


class PicoProtocol(websockets.WebSocketServerProtocol):

    def http_unauthorized(self, message, realm='pico'):
        return (
            http.HTTPStatus.UNAUTHORIZED,
            [("WWW-Authenticate", build_www_authenticate_basic(realm))],
            message.encode(),
        )

    def serve_file(self, path):
        document = Path(__file__, '..', Path(path).name).resolve()
        if not document.is_file():
            document = Path(__file__, '..', 'pico.html').resolve()
            content_type = 'text/html; charset=utf-8'
        elif path.endswith('.js'):
            content_type = 'application/javascript; charset=utf-8'
        elif path.endswith('.css'):
            content_type = 'text/css; charset=utf-8'
        else:
            content_type = 'text/plain; charset=utf-8'
        return (
            http.HTTPStatus.OK,
            [('Content-Type', content_type)],
            document.read_bytes(),
        )

    async def process_request(self, path, request_headers):
        if request_headers.get('Upgrade') != 'websocket':
            return self.serve_file(path)
        try:
            authorization = request_headers['Authorization']
            self.username, password = parse_authorization_basic(authorization)
        except Exception as error:
            return self.http_unauthorized('No username found')
        return await super().process_request(path, request_headers)


def get_usernames(sockets):
    return sorted(ws.username for ws in sockets)


async def send_json_many(targets, **data):
    for websocket in list(targets):
        await send_json(websocket, **data)


async def send_json(websocket, **data):
    try:
        await websocket.send(json.dumps(data))
    except websockets.exceptions.ConnectionClosed:
        pass


async def recv_json(websocket):
    try:
        return json.loads(await websocket.recv())
    except websockets.exceptions.ConnectionClosed:
        return {'action': 'logout'}
    except json.decoder.JSONDecodeError:
        return {}


async def core(ws, path, server_name):
    sockets = rooms[path]

    while True:
        if ws in sockets:
            data = await recv_json(ws)
        else:
            data = {'action': 'login'}

        ts = datetime.datetime.now().isoformat()
        reply = functools.partial(send_json, websocket=ws, ts=ts)
        error = functools.partial(reply, kind='error')
        broadcast = functools.partial(send_json_many, targets=set(sockets), ts=ts)

        if 'action' not in data:
            await error(info='Message without action is invalid')

        elif data['action'] == 'login':
            if not ws.username:
                await error(info='Username not allowed')
                break
            if ws.username in get_usernames(sockets):
                await error(info='Username taken')
                break

            sockets.add(ws)
            online = get_usernames(sockets)
            await reply(kind='update', users=online, info=f'Welcome to {path}', username=ws.username)
            await broadcast(kind='update', users=online, info=f'{ws.username} joined')

        elif data['action'] == 'post':
            text = data['text']
            if not text:
                continue
            await broadcast(kind='post', source=ws.username, text=text)

        elif data['action'] == 'logout':
            sockets.discard(ws)
            await broadcast(kind='update', users=get_usernames(sockets), info=f'{ws.username} left')
            break


async def start_server(host, port, server_name):
    bound_core = functools.partial(core, server_name=server_name)
    return await websockets.serve(bound_core, host, port, create_protocol=PicoProtocol)


if __name__ == '__main__':
    host, port = 'localhost', 9753
    loop = asyncio.get_event_loop()
    loop.run_until_complete(start_server(host, port, 'PicoChat'))
    print(f'Running on {host}:{port}')
    loop.run_forever()