An incomplete WebSocket client based on only socket, ssl, and uuid in Python

A few days ago, I couldn't get a WebSocket library working with another library on Python 3.10. So to avoid those dependencies, I implemented my WebSocket client on a low-level socket API. I implemented one in Common Lisp first. Then I translated it to Python. My WebSocket client is very far away from being completed, but at least it can run.

The first part is import. I imported only three things that are socket, ssl, and uuid.

import socket
import ssl
import uuid

The second part is based on HTTP headers for telling the server to switch to the WebSocket mode.

def upgrade(s, conn_info):
    with s.makefile(mode = 'rw', encoding = "ISO-8859-1") as f:
        f.write(f'GET {conn_info["path"]} HTTP/1.1\r\n')
        f.write(f'Host: {conn_info["host"]}\r\n')
        f.write("Connection: Upgrade\r\n")
        f.write("Upgrade: websocket\r\n")
        f.write(f'Sec-Websocket-Key: {str(uuid.uuid1())}\r\n')
        f.write("Sec-WebSocket-Version: 13\r\n")
        f.write("\r\n")
        f.flush()
        
        # reading response
        for line in f:
            if line == "\n":
                break

The third part is reading a content length. I assume that the server keep sending text contents. In read_payload_len(), s.recv(1) read a byte from the socket s. & 0x7F is for masking only 7 bits. If length is 127, the length is in the next 4 bytes instead. If length is 126, the length is in the next 2 bytes. Otherwise the function just returns the length.

def read_payload_len(s):
    match s.recv(1)[0] & 0x7F:
        case 127:
            return read_extra_len(s, 4)
        case 126:
            return read_extra_len(s, 2)
        case l:
            return l

The read_extra_len reads bytes and turn them to integer.

def read_extra_len(s, num_of_bytes):
    buf = s.recv(num_of_bytes)
    len = 0
    for i in range(num_of_bytes):
        len += buf[i] << (8 * (num_of_bytes - i - 1))
    return len

In the fourth part, we read a web socket frame. The program determine a frame type from opcode, which in the last 4 bits of the header. I should implement PING-PONG part but I didn't. According to RFC6455, which I forgot to mention before, opcode == 0x1 means the frame is a text frame. So the program reads payload length and reads the payload.

def read_frame(s):
    header0 = s.recv(1)[0]
    opcode = header0 & 0x0F
    match opcode:
        case 0x1:
            payload_len = read_payload_len(s)
            print(s.recv(payload_len))
        case 0x9:
            print("PING")
        case 0xA:
            print("PONG")

The last part is for opening connections and SSL/TLS wrapper. The function created a socket and wrapped it with TLS/SSL wrapper. The sending upgrade message to ask the server to switch to Websocket mode and the keep reading frames.

def connect(conn_info):
    ctx = ssl.create_default_context()
    with socket.create_connection((conn_info["host"], conn_info["port"])) as s:
        with ctx.wrap_socket(s, server_hostname = conn_info["host"]) as ss:
            upgrade(ss, conn_info)
            while True:
                read_frame(ss)

In the final part, I cannot find any public Websocket endpoint besides ones from cryptocurrency exchanges. So I put Bitkub API.

connect({"host": "api.bitkub.com",
         "port": 443,
         "path": "/websocket-api/market.trade.thb_btc"})

And it works, but if you are unlucky, you will get PING instead.

> python3.10 http_ex.py 
b'{"amt":0.00004661,"bid":86911896,"rat":2140000,"sid":82187323,"stream":"market.trade.thb_btc","sym":"THB_BTC","ts":1636729677,"txn":"BTCSELL0011078574"}\n{"amt":0.23306074,"bid":86912246,"rat":2140000,"sid":82187325,"stream":"market.trade.thb_btc","sym":"THB_BTC","ts":1636729677,"txn":"BTCSELL0011078576"}'
b'{"amt":0.00466121,"bid":86912014,"rat":2140000,"sid":82187324,"stream":"market.trade.thb_btc","sym":"THB_BTC","ts":1636729677,"txn":"BTCSELL0011078575"}'