#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