js-web-socket-client.html

<meta name="viewport" content="width=device-width">
<script src="//ajax.aspnetcdn.com/ajax/jQuery/jquery-1.11.0.js"></script>
<style>
#in { width: 500px; }
</style>
<input type="text" id="in">
<div id="out"></div>
<script>
sock = new WebSocket('ws://baro.sl:7942/wakana');
//sock = new WebSocket('ws://baro.sl:8765/wakana');
sock.onopen = function() {
    $('<p>').html('onopen').appendTo('#out');

    sock.send('Hello. This is a client.');
//    sock.send(new Array(102400+1).join('A'));
};
sock.onerror = function() {
    $('<p>').html('onerror').appendTo('#out');
};
sock.onclose = function() {
    $('<p>').html('onclose').appendTo('#out');
};
sock.onmessage = function(e) {
    $('<p>').html('onmessage('+e.data.length+'): '+e.data).appendTo('#out');
};

$('#in').keydown(function(ev) {
    if (ev.keyCode == 13) {
        var text = $(ev.target).val();
        $(ev.target).val('');

        if (text.length) {
            $('<p>').html('send('+text.length+'): '+text).appendTo('#out');
            sock.send(text);
        }
    }
});

$(function() {
    $('#in').focus();
});
</script>

python-web-socket-server.py

#!/usr/bin/env python3

import socket
import re
import base64
import hashlib
import struct
import threading

def sock_recv(sock, size):
    data = b''
    while len(data) != size:
        buf = sock.recv(size-len(data))
        if not buf: raise Exception('Socket buffer underrun')
        data += buf
    return data

def ws_recv(sock, proto=0):
    if proto:
        buf = sock.recv(1)
        if not buf: return b''
        assert buf[0] == 0x0
        data = b''
        while True:
            buf = sock.recv(1)
            if not buf: raise Exception('Unexpected end of frame')
            if buf[0] == 0xff: break
            data += buf
        return data

    buf = sock_recv(sock, 2)
    fin = buf[0] >> 7
    op = buf[0] & 0xf
    mask = buf[1] >> 7
    plen = buf[1] & 0x7f

#    print('op={} fin={} plen={} mask={}'.format(op, fin, plen, mask))
    assert fin == 1 # TODO: No fragmentation support

    data = b''

    if op < 3:
        if plen == 126:
            buf = sock_recv(sock, 2)
            plen, = struct.unpack('>H', buf)
        elif plen == 127:
            buf = sock_recv(sock, 8)
            plen, = struct.unpack('>Q', buf)

        if mask: mask_key = sock_recv(sock, 4)

        data = sock_recv(sock, plen)

        if mask: data = bytes(data[i] ^ mask_key[i%4] for i in range(len(data)))

        return data

    elif op == 0x8:
        print('connection close')
#        sock.close()
        return ''

    else:
        raise Exception('Unhandled frame: op={}'.format(op))

def ws_send(sock, buf, proto=0):
    if proto:
        buf = b'\x00' + buf + b'\xff'
    elif len(buf) < 126:
        buf = bytes((0x81, len(buf))) + buf
    elif len(buf) < 0x10000:
        buf = bytes((0x81, 126)) + struct.pack('>H', len(buf)) + buf
    elif len(buf) < 0x8000000000000000:
        buf = bytes((0x81, 127)) + struct.pack('>Q', len(buf)) + buf
    else:
        raise Exception('Frame too big')

    sock.send(buf)

def handshake(sock):
    header = b''
    serv_key = b''

    while True:
        buf = sock.recv(1024)
        if not buf: raise Exception('Unable to handshake')

        header += buf
        try: header, buf = header.split(b'\r\n\r\n', 1)
        except ValueError: continue

#        print(header.decode('utf-8'))
#        print('----')

        try:
            cli_key = re.search(br'(?im)^Sec-WebSocket-Key: (\S+)', header).group(1)
            serv_key = base64.b64encode(hashlib.sha1(cli_key+b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11').digest())

            resp = 'HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n\r\n'.format(serv_key.decode('utf-8')).encode('utf-8')
            proto = 0

        except AttributeError:
            assert len(buf) == 8 # FIXME: Change this to use sock_recv

            cli_key1 = re.search(br'(?im)^Sec-WebSocket-Key1: ([^\r\n]+)', header).group(1)
            cli_key2 = re.search(br'(?im)^Sec-WebSocket-Key2: ([^\r\n]+)', header).group(1)
            cli_key3 = buf

            digit1 = int(b''.join(re.findall(b'[0-9]', cli_key1)))
            digit2 = int(b''.join(re.findall(b'[0-9]', cli_key2)))
            spaces1 = cli_key1.count(b' ')
            spaces2 = cli_key2.count(b' ')
            num_key1 = digit1//spaces1
            num_key2 = digit2//spaces2

            serv_key = hashlib.md5(struct.pack('>L', num_key1)+struct.pack('>L', num_key2)+cli_key3).digest()

            path = re.search(br'(?i)^GET (\S+) HTTP/', header).group(1).decode('utf-8')
            origin = re.search(br'(?im)^Origin: ([^\r\n]+)', header).group(1).decode('utf-8')
            loc = 'ws://{}{}'.format(re.search(br'(?im)^Host: ([^\r\n]+)', header).group(1).decode('utf-8'), path)

            resp = 'HTTP/1.1 101 WebSocket Protocol Handshake\r\nUpgrade: WebSocket\r\nConnection: Upgrade\r\nSec-WebSocket-Origin: {origin}\r\nSec-WebSocket-Location: {loc}\r\n\r\n'.format(origin=origin, loc=loc).encode('utf-8') + serv_key
            proto = 1

        sock.send(resp)

#        print(resp)
#        print('----')

        return proto

def proc(sock):
    proto = handshake(sock)
    print('* handshake: proto={}'.format(proto))

    ws_send(sock, b'Hi. This is a server.', proto)

    while True:
        buf = ws_recv(sock, proto)
        if not buf: break

        msg = buf.decode('utf-8')
        print('* text({}): {}'.format(len(msg), msg))

        ws_send(sock, msg.encode('utf-8'), proto)

def main():
    serv_sock = socket.socket()
    serv_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    serv_sock.bind(('', 7942))
    serv_sock.listen(5)

    while True:
        sock, addr = serv_sock.accept()
        print('* new connection: {}'.format(addr))

        thrd = threading.Thread(target=proc, args=(sock,))
        thrd.daemon = True
        thrd.start()

if __name__ == '__main__':
    main()