mirror of
https://github.com/hak5/nano-tetra-modules.git
synced 2025-10-29 16:58:09 +00:00
51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
import struct
|
|
import SocketServer
|
|
import time
|
|
from base64 import b64encode
|
|
from hashlib import sha1
|
|
from mimetools import Message
|
|
from StringIO import StringIO
|
|
|
|
# This class is taken and modified from https://gist.github.com/jkp/3136208
|
|
class WebSocketsHandler(SocketServer.StreamRequestHandler):
|
|
magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
|
|
|
|
def setup(self):
|
|
SocketServer.StreamRequestHandler.setup(self)
|
|
|
|
self.server.websockets.append(self)
|
|
self.handshake_done = False
|
|
|
|
def handle(self):
|
|
while self.server.running:
|
|
if not self.handshake_done:
|
|
self.handshake()
|
|
else:
|
|
time.sleep(5)
|
|
|
|
def send_message(self, message):
|
|
self.request.send(chr(129))
|
|
length = len(message)
|
|
if length <= 125:
|
|
self.request.send(chr(length))
|
|
elif length >= 126 and length <= 65535:
|
|
self.request.send(chr(126))
|
|
self.request.send(struct.pack(">H", length))
|
|
else:
|
|
self.request.send(chr(127))
|
|
self.request.send(struct.pack(">Q", length))
|
|
self.request.send(message)
|
|
|
|
def handshake(self):
|
|
data = self.request.recv(1024).strip()
|
|
headers = Message(StringIO(data.split('\r\n', 1)[1]))
|
|
if headers.get("Upgrade", None) != "websocket":
|
|
return
|
|
key = headers['Sec-WebSocket-Key']
|
|
digest = b64encode(sha1(key + self.magic).hexdigest().decode('hex'))
|
|
response = 'HTTP/1.1 101 Switching Protocols\r\n'
|
|
response += 'Upgrade: websocket\r\n'
|
|
response += 'Connection: Upgrade\r\n'
|
|
response += 'Sec-WebSocket-Accept: %s\r\n\r\n' % digest
|
|
self.handshake_done = self.request.send(response)
|