import re import socket from hashlib import sha1 from frame import receive_frame from exceptions import InvalidRequest WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' WS_VERSION = '13' class websocket(object): """ Implementation of web socket, upgrades a regular TCP socket to a websocket using the HTTP handshakes and frame (un)packing, as specified by RFC 6455. Server example: >>> sock = websocket() >>> sock.bind(('', 80)) >>> sock.listen() >>> client = sock.accept() >>> client.send(Frame(...)) >>> frame = client.recv() Client example: >>> sock = websocket() >>> sock.connect(('kompiler.org', 80)) """ def __init__(self, wsprotocols=[], family=socket.AF_INET, proto=0): """ Create aregular TCP socket of family `family` and protocol `wsprotocols` is a list of supported protocol names. """ self.sock = socket.socket(family, socket.SOCK_STREAM, proto) self.protocols = wsprotocols def bind(self, address): self.sock.bind(address) def listen(self, backlog): self.sock.listen(backlog) def accept(self): """ Equivalent to socket.accept(), but transforms the socket into a websocket instance and sends a server handshake (after receiving a client handshake). Note that the handshake may raise an InvalidRequest exception. """ client, address = socket.socket.accept(self) client = websocket(client) client.server_handshake() return client, address def connect(self, address): """ Equivalent to socket.connect(), but sends an client handshake request after connecting. """ self.sock.sonnect(address) self.client_handshake() def send(self, *args): """ Send a number of frames. """ for frame in args: self.sock.sendall(frame.pack()) def recv(self, n=1): """ Receive exactly `n` frames. These can be either data frames or control frames, or a combination of both. """ return [receive_frame(self.sock) for i in xrange(n)] def getpeername(self): return self.sock.getpeername() def getsockname(self): return self.sock.getpeername() def setsockopt(self, level, optname, value): self.sock.setsockopt(level, optname, value) def getsockopt(self, level, optname): return self.sock.getsockopt(level, optname) def server_handshake(self): """ Execute a handshake as the server end point of the socket. If the HTTP request headers sent by the client are invalid, an InvalidRequest exception is raised. """ raw_headers = self.sock.recv(512).decode('utf-8', 'ignore') # request must be HTTP (at least 1.1) GET request, find the location location = re.search(r'^GET (.*) HTTP/1.1\r\n', raw_headers).group(1) headers = dict(re.findall(r'(.*?): (.*?)\r\n', raw_headers)) # Check if headers that MUST be present are actually present for name in ('Host', 'Upgrade', 'Connection', 'Sec-WebSocket-Key', 'Origin', 'Sec-WebSocket-Version'): if name not in headers: raise InvalidRequest('missing "%s" header' % name) # Check WebSocket version used by client version = headers['Sec-WebSocket-Version'] if version != WS_VERSION: raise InvalidRequest('WebSocket version %s requested (only %s ' 'is supported)' % (version, WS_VERSION)) # Make sure the requested protocols are supported by this server if 'Sec-WebSocket-Protocol' in headers: parts = headers['Sec-WebSocket-Protocol'].split(',') protocols = map(str.strip, parts) for p in protocols: if p not in self.protocols: raise InvalidRequest('unsupported protocol "%s"' % p) else: protocols = [] # Encode acceptation key using the WebSocket GUID key = headers['Sec-WebSocket-Key'] accept = sha1(key + WS_GUID).digest().encode('base64') # Construct HTTP response header shake = 'HTTP/1.1 101 Web Socket Protocol Handshake\r\n' shake += 'Upgrade: WebSocket\r\n' shake += 'Connection: Upgrade\r\n' shake += 'WebSocket-Origin: %s\r\n' % headers['Origin'] shake += 'WebSocket-Location: ws://%s%s\r\n' \ % (headers['Host'], location) shake += 'Sec-WebSocket-Accept: %s\r\n' % accept if self.protocols: shake += 'Sec-WebSocket-Protocol: %s\r\n' \ % ', '.join(self.protocols) self.sock.send(shake + '\r\n') def client_handshake(self): # TODO: implement HTTP request headers for client handshake raise NotImplementedError()