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.
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.