WebSocket Echo Server using asyncio

websock_echo_server.py

#!/usr/bin/env python3

from struct import pack, unpack
from base64 import b64encode
from hashlib import sha1

import asyncio

async def _ws_handle_handshake(reader, writer):
  data = await reader.read(1024)
  data = data.decode().strip()

  headers = [h.split(':', 1) for h in data.split('\r\n')[1:]]
  headers = {k.strip().lower(): v.strip() for k, v in headers}
  if headers.get("upgrade", None) != "websocket":
    return False

  key = headers['sec-websocket-key'].encode()
  WEBSOCKET_MAGIC_STR = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
  digest = b64encode(sha1(key + WEBSOCKET_MAGIC_STR).digest())

  response  = 'HTTP/1.1 101 Switching Protocols\r\n'
  response += 'Upgrade: websocket\r\n'
  response += 'Connection: Upgrade\r\n'
  response += 'Sec-WebSocket-Accept: %s\r\n\r\n' % digest.decode()

  writer.write(response.encode())
  await writer.drain()
  return True

async def _ws_handle_read(reader):
  data = await reader.read(2)
  if not data: return 8, None

  opcode = data[0] & 15
  length = data[1] & 127
  if length == 126:
    data = await reader.read(2)
    length = unpack(">H", data)[0]
  elif length == 127:
    data = await reader.read(8)
    length = unpack(">Q", data)[0]

  masks = await reader.read(4)
  payload = await reader.read(length)
  return opcode, ''.join(chr(b ^ masks[i % 4]) for i, b in enumerate(payload))

async def _ws_handle_write(writer, message):
  writer.write(bytes([129]))
  length = len(message)
  if length <= 125:
    writer.write(bytes([length]))
  elif length >= 126 and length <= 65535:
    writer.write(bytes([126]))
    writer.write(pack(">H", length))
  else:
    writer.write(bytes([127]))
    writer.write(pack(">Q", length))
  writer.write(message)
  await writer.drain()

async def handle_ws_echo(reader, writer):
  if await _ws_handle_handshake(reader, writer):
    while True:
      opcode, message = await _ws_handle_read(reader)
      if opcode == 8 or message is None: break

      if opcode == 1:
        print(f"Received Text: {message!r}")
      else:
        print(f"Received {opcode} {message!r}")

      await _ws_handle_write(writer, ('echo: ' + message).encode())

  print("Close the connection")
  writer.close()

async def main():
  server = await asyncio.start_server(handle_ws_echo, '127.0.0.1', 57025)

  addr = server.sockets[0].getsockname()
  print(f'Serving on {addr}')

  async with server:
    await server.serve_forever()

if __name__ == '__main__':
  asyncio.run(main())