Welcome, guest | Sign In | My Account | Store | Cart

This tries its best to be a replacement for the regular socket module.

It supports only sending and receiving but should be useful enough.

The only real difference should be that you can't specify the number of bytes is received, instead do

for message in socket.recv():
    print(message)

Revision 2: Added proper message receiving. Previously it just requested a ton of data. Now it reads 2 bytes, determines the length, then requests that much.

Python, 138 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
#WebSocket module, designed to be exactly like socket

import socket
import struct

from base64 import b64encode
import hashlib

import json

AF_INET = 0
WSOCK_STREAM = 1

class error(Exception):
    pass


class websocketclient:
    def __init__(self, socket):
        self.socket = socket
        #self.socket.setblocking(False)

    def __decodeWebSocket (self, data):  #JS PORT
        if data == b'\x81': return
        datalength = data[1] & 127;
        indexFirstMask = 2;
        if (datalength == 126):
            indexFirstMask = 4;
        elif (datalength == 127):
            indexFirstMask = 10;
        
        masks = data[indexFirstMask:indexFirstMask + 4]
        i = indexFirstMask + 4;
        index = 0;
        output = "";
        while (i < len(data)):
            output += chr(data[i] ^ masks[index % 4]);
            i += 1
            index += 1
    
        return output;

    def __encodeWebSocket(self, bytesRaw):
        print(bytesRaw)

        bytesFormatted = []
        bytesFormatted.append(struct.pack('B', 129))
        if (len(bytesRaw) <= 125):
            bytesFormatted.append(struct.pack('B', len(bytesRaw)))
        elif (len(bytesRaw) >= 126 and len(bytesRaw) <= 65535):
            bytesFormatted.append(struct.pack('B', 126));
            bytesFormatted.append(struct.pack('B', ( len(bytesRaw) >> 8 ) & 255));
            bytesFormatted.append(struct.pack('B', ( len(bytesRaw)      ) & 255));
            
        else:
            bytesFormatted.append(struct.pack('B', 127));
            bytesFormatted.append(struct.pack('B', ( len(bytesRaw) >> 56 ) & 255));
            bytesFormatted.append(struct.pack('B', ( len(bytesRaw) >> 48 ) & 255));
            bytesFormatted.append(struct.pack('B', (len( bytesRaw) >> 40 ) & 255));
            bytesFormatted.append(struct.pack('B', (len( bytesRaw) >> 32 ) & 255));
            bytesFormatted.append(struct.pack('B', (len( bytesRaw) >> 24 ) & 255));
            bytesFormatted.append(struct.pack('B', (len( bytesRaw) >> 16 ) & 255));
            bytesFormatted.append(struct.pack('B', ( len(bytesRaw) >>  8 ) & 255));
            bytesFormatted.append(struct.pack('B', ( len(bytesRaw)       ) & 255));
        
        for i in range(len(bytesRaw)):
            bytesFormatted.append(struct.pack('B', ord(bytesRaw[i])))
        
        return bytesFormatted;

    def recv(self):
        length, = struct.unpack('xB', self.socket.recv(2))
        firstbyte = length
        length &= 127

        addbytes = b""

        if length == 126:
            #Requires 2 more bytes of data
            addbytes = self.socket.recv(2)
            length, = struct.unpack("!H", addbytes)

        if length == 127:
            addbytes = self.socket.recv(8)
            length, = struct.unpack("!Q", addbytes)

        length += 4

        return [self.__decodeWebSocket(b'\x81' + struct.pack('B', firstbyte) + addbytes + self.socket.recv(length))]


        
        '''retlist = []
        for chunk in buffer.split(b'\x81'):
            pop = self.__decodeWebSocket(b'\x81' + chunk)
            if pop: retlist.append(pop)

        return retlist'''

    def send(self, string):
        self.socket.send(b''.join(self.__encodeWebSocket(string)))


class websocket:
    def __init__(self, flag1, flag2):
        if (flag1, flag2) != (AF_INET, WSOCK_STREAM):
            raise (error, "Only supports flags AF_INET, WSOCK_STREAM!")

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    def bind(self, data):
        self.socket.bind(data)

    def listen(self, conns):
        self.socket.listen(conns * 5)

    def accept(self):
        client, addr = self.socket.accept()
        handshake = client.recv(10000).decode().split("\r\n")

        newshake = {}
        for item in handshake:
            if ":" in item:
                newshake[item.split(":")[0]] = item.split(":")[1][1:]

        handshake = newshake

        m = hashlib.sha1()

        thehash = handshake['Sec-WebSocket-Key'] + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
        m.update(thehash.encode())
        thehash = b64encode(m.digest())

        header = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + thehash + b"\r\n\r\n"           
        client.send(header)
        
        return websocketclient(client), addr

Probably this needs more error checking, but it's easily add-able.