def kill(self, code=1000, error=b''): """Force to close the connection Keyword Arguments: code {int} -- Closing code according to RFC. (default: {1000}) error {bytes} -- Error message to append on closing frame. (default: {''}) """ logging.websocket('--- KILL (CONTROLLER) ---', repr(self._WSClient.conn)) if not self._WSClient.hasStatus('CLOSED'): _WSEncoder = WSEncoder() data = struct.pack('!H', code) if len(error): data += error logging.websocket('--- KILL FRAME ---') logging.websocket('Code:', code) logging.websocket('Error:', error) logging.websocket(repr(self._WSClient.conn)) try: bytes = _WSEncoder.close(data) except ValueError as error: self._WSClient.close() else: self._WSClient.send(bytes) self._WSClient.close()
def stop(self): """Stop all clients """ self.listening = False while len(self.clients): self.clients.pop()._WSController.kill() self.s.close() logging.websocket('--- THAT\'S ALL FOLKS ---')
def close(self): """Close connection """ if self._WSServer._WSHandler is not None: self._WSServer._WSHandler.onClose(self.conn) logging.websocket(repr(self.conn)) if not self.hasStatus('CLOSED'): self.setStatus('CLOSED') self.conn.close()
def ping(self): """Send a ping """ logging.websocket('--- PING (CONTROLLER) ---') if self._WSClient.hasStatus('OPEN'): _WSEncoder = WSEncoder() try: bytes = _WSEncoder.ping(b'Application data') except ValueError as error: self._WSClient._WSServer.remove(self._WSClient) self.kill(1011, 'WSEncoder error: ' + str(error)) else: self._WSClient.send(bytes)
def receive(self, bufsize): """Real socket bytes reception Arguments: bufsize {int} -- Buffer size to return """ bytes = self.conn.recv(bufsize) # Receive byte string if not bytes: logging.websocket('Client left', repr(self.conn)) self._WSServer.remove(self) self.close() return b'' return bytes
def run(self): """Start server Keyword Arguments: host {str} -- WebSocket server or ip to join. (default: {'localhost'}) port {int} -- Port to join. (default: {9999}) maxclients {int} -- Max clients which can connect at the same time. (default: {20}) """ self.s = socket.socket() self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) self.s.bind((self.host, self.port)) self.s.listen(1) self.listening = True while self.listening: conn, addr = self.s.accept() logging.websocket('New client host/address:', addr) if len(self.clients) == self.maxclients: logging.websocket('Too much clients - connection refused:', repr(conn)) conn.close() else: _WSClient = WSClient(self) self.clients.append(_WSClient) logging.websocket('Total clients:', len(self.clients)) threading.Thread(target=_WSClient.handle, args=(conn, addr)).start()
def send(self, bytes): """Send a unicast frame Arguments: bytes {bytes} -- Bytes to send """ if not self.hasStatus('CLOSED'): logging.websocket('--- SEND UNICAST ---') logging.websocket(repr(self.conn)) logging.websocket(repr(bytes), '[', str(len(bytes)), ']') if self._WSServer._WSHandler is not None: self._WSServer._WSHandler.onSend(bytes) lock = threading.Lock() lock.acquire() self.conn.send(bytes) lock.release() logging.websocket('--- END UNICAST ---')
def handle(self, conn, addr): """Handle incoming datas Arguments: conn {socket} -- Socket of WebSocket client addr {address} -- Adress of WebSocket client """ self.conn = conn self.addr = addr self.setStatus('CONNECTING') try: getRequest = self.handshake() except ValueError as error: self._WSServer.remove(self) self.close() raise ValueError('Client rejected: ' + str(error)) else: _WSDecoder = WSDecoder() self.setStatus('OPEN') if self._WSServer._WSHandler is not None: self._WSServer._WSHandler.onConnect(self.conn, getRequest) while self.hasStatus('OPEN'): try: ctrl, data = _WSDecoder.decode(self) except ValueError as e: closing_code, message = (1000, e) parts = str(e).split('|') if len(parts) == 2: closing_code = int(parts[0]) message = parts[1] if self.hasStatus('OPEN'): self._WSController.kill(closing_code, ('WSDecoder::' + str(message)).encode('UTF-8')) break else: logging.websocket('--- INCOMING DATAS ---') self._WSController.run(ctrl, data)
def send(self, bytes): """Send a multicast frame Arguments: bytes {bytes} -- Bytes to send """ logging.websocket('--- SEND MULTICAST ---') logging.websocket(repr(bytes)) for _WSClient in self.clients: _WSClient.send(bytes) logging.websocket('multicast send finished')
def run(self, ctrl, data): """Handle incoming datas Arguments: ctrl {ctrl} -- Control dictionnary for data data {string} -- Decoded data, text or binary """ logging.websocket('--- CONTROLLER ---') logging.websocket(repr(self._WSClient.conn)) _WSEncoder = WSEncoder() # CONTROLS if ctrl['opcode'] == 0x9: # PING logging.websocket('--- PING FRAME --- ') logging.websocket(repr(self._WSClient.conn)) try: bytes = _WSEncoder.pong('Application data') except ValueError as error: self._WSClient._WSServer.remove(self._WSClient) self.kill(1011, 'WSEncoder error: ' + str(error)) else: self._WSClient.send(bytes) if ctrl['opcode'] == 0xA: # PONG logging.websocket('--- PONG FRAME ---') logging.websocket(repr(self._WSClient.conn)) if len(data): logging.websocket('Pong frame datas:', str(data)) if ctrl['opcode'] == 0x8: # CLOSE logging.websocket('--- CLOSE FRAME ---') logging.websocket(repr(self._WSClient.conn)) self._WSClient._WSServer.remove(self._WSClient) # closing was initiated by server if self._WSClient.hasStatus('CLOSING'): self._WSClient.close() # closing was initiated by client if self._WSClient.hasStatus('OPEN'): self._WSClient.setStatus('CLOSING') self.kill(1000, b'Goodbye !') # the two first bytes MUST contains the exit code, follow optionnaly with text data not shown to clients if len(data) >= 2: code, data = self.array_shift(data, 2) status = '' if code in WSSettings.CLOSING_CODES: logging.websocket('Closing frame code:', code) if len(data): logging.websocket('Closing frame data:', data) # DATAS if ctrl['opcode'] == 0x1: # TEXT logging.websocket('--- TEXT FRAME ---') logging.websocket(repr(self._WSClient.conn)) if len(data): try: # Handle message if possible if self._WSClient._WSServer._WSHandler is not None: self._WSClient._WSServer._WSHandler.onMessage( data, self._WSClient) except ValueError as error: self._WSClient._WSServer.remove(self._WSClient) self.kill(1011, ('WSEncoder error: ' + str(error)).encode('UTF-8')) if ctrl['opcode'] == 0x0: # CONTINUATION logging.websocket('--- CONTINUATION FRAME ---', repr(self._WSClient.conn)) pass if ctrl['opcode'] == 0x2: # BINARY logging.websocket('--- BINARY FRAME ---', repr(self._WSClient.conn)) pass
def decode(self, _WSClient): """Decode on the fly data from WSClient Arguments: _WSClient {WSClient} -- WebSocket Client 0 1 2 3 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 +-+-+-+-+-------+-+-------------+-------------------------------+ |F|R|R|R| opcode|M| Payload len | Extended payload length | |I|S|S|S| (4) |A| (7) | (16/64) | |N|V|V|V| |S| | (if payload len==126/127) | | |1|2|3| |K| | | +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - + | Extended payload length continued, if payload len == 127 | + - - - - - - - - - - - - - - - +-------------------------------+ | |Masking-key, if MASK set to 1 | +-------------------------------+-------------------------------+ | Masking-key (continued) | Payload Data | +-------------------------------- - - - - - - - - - - - - - - - + : Payload Data continued ... : + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + | Payload Data continued ... | OPCODES USED FOR ERRORS: 1002: PROTOCOL_ERROR 1003: UNSUPPORTED_DATA_TYPE 1007: INVALID_PAYLOAD 1011: UNEXPECTED_CONDITION_ENCOUTERED_ON_SERVER """ # decode first byte b = _WSClient.read(1) if not len(b): raise ValueError('1011|Reading first byte failed.') b1 = ord(b) fin = b1 >> 7 & 1 rsv1 = b1 >> 6 & 1 rsv2 = b1 >> 5 & 1 rsv3 = b1 >> 4 & 1 opcode = b1 & 0xf if not opcode in WSSettings.OPCODES: raise ValueError('1002|Wrong opcode.') # decode second byte b = _WSClient.read(1) if not len(b): raise ValueError('1011|Reading second byte failed.') b2 = ord(b) mask = b2 >> 7 & 1 if mask != 0x1: raise ValueError('1002|Client datas MUST be masked.') # decode data length (without mask size) length = b2 & 0x7f # RFC : If length is 126, the following 2 bytes must be interpreted as a 16-bit unsigned integer # are the payload length if length == 0x7e: b = _WSClient.read(2) if not len(b): raise ValueError('1011|Reading length failed.') length_bytes = b length = struct.unpack("!H", length_bytes)[0] # RFC : If length is 127, the following 8 bytes must be interpreted as a 64-bit unsigned integer (the # most significant bit MUST be 0) are the payload length elif length == 0x7f: b = _WSClient.read(8) if not len(b): raise ValueError('1011|Reading length failed.') length_bytes = b length = struct.unpack("!Q", length_bytes)[0] # decode mask key mask_key = _WSClient.read(4) if not len(mask_key): raise ValueError('1011|Reading mask key failed.') # Note: we are trying here to read exactly the data amount, so we don't need MESSAGE_TOO_BIG data = _WSClient.read(length) if not len(data): raise ValueError('1011|Reading data failed.') data = self.unmask(mask_key, data) if opcode == 0x1: try: data = data.decode('utf-8') except UnicodeError: raise ValueError( '1003|Client text datas MUST be UTF-8 encoded.') ctrl = { 'fin': fin, 'opcode': opcode, 'rsv1': rsv1, 'rsv2': rsv2, 'rsv3': rsv3, } logging.websocket('After decode:', repr(ctrl), repr(data)) return ctrl, data
def remove(self, _WSClient): if _WSClient in self.clients: logging.websocket('Client left:', repr(_WSClient.conn)) self.clients.remove(_WSClient)
def handshake(self): """Send handshake according to RFC """ headers = {} # Ignore first line with GET getRequest = self.readlineheader() while self.hasStatus('CONNECTING'): if len(headers) > 64: raise ValueError('Header too long.') line = self.readlineheader() if not self.hasStatus('CONNECTING'): raise ValueError('Client left.') if len(line) == 0 or len(line) == 1024: raise ValueError('Invalid line in header.') if line == '\r\n': break # Take care with strip ! # >>> import string;string.whitespace # '\t\n\x0b\x0c\r' line = line.strip() # Take care with split ! # >>> a='key1:value1:key2:value2';a.split(':',1) # ['key1', 'value1:key2:value2'] kv = line.split(':', 1) if len(kv) == 2: key, value = kv k = key.strip().lower() v = value.strip() headers[k] = v else: raise ValueError('Invalid header key/value.') if not len(headers): raise ValueError('Reading headers failed.') if not 'sec-websocket-version' in headers: raise ValueError('Missing parameter "Sec-WebSocket-Version".') if not 'sec-websocket-key' in headers: raise ValueError('Missing parameter "Sec-WebSocket-Key".') if not 'host' in headers: raise ValueError('Missing parameter "Host".') if not 'origin' in headers: raise ValueError('Missing parameter "Origin".') if int(headers['sec-websocket-version']) != WSSettings.VERSION: raise ValueError('Wrong protocol version %s.' % WSSettings.VERSION) logging.websocket('--- RECEIVED HEADERS ---') for c in headers: logging.websocket (c + ':', headers[c]) logging.websocket('------------------------') accept = base64.b64encode(hashlib.sha1(headers['sec-websocket-key'].encode('UTF-8') + b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11').digest()).decode('UTF-8') bytes = ('HTTP/1.1 101 Switching Protocols\r\n' 'Upgrade: websocket\r\n' 'Connection: Upgrade\r\n' 'Sec-WebSocket-Origin: %s\r\n' 'Sec-WebSocket-Location: ws://%s\r\n' 'Sec-WebSocket-Accept: %s\r\n' 'Sec-WebSocket-Version: %s\r\n' '\r\n') % (headers['origin'], headers['host'], accept, headers['sec-websocket-version']) logging.websocket('--- HANDSHAKE ---') logging.websocket(bytes) logging.websocket('-----------------') self.send(bytes.encode('UTF-8')) return getRequest
def encode(self, opcode=0x1, data='', fin=1, mask=1, rsv1=0, rsv2=0, rsv3=0): """Encoding function for all types Keyword Arguments: opcode {hex} -- Operation code according to RFC (default: {0x1}) data {str} -- UTF-8 text to send (default: {''}) fin {bit} -- Bit which define if the frame is the last one (default: {1}) mask {bit} -- Bit which define if datas must be masked or not (default: {1}) rsv1 {bit} -- Reserved bit for future usage. Do not use. (default: {0}) rsv2 {bit} -- Reserved bit for future usage. Do not use. (default: {0}) rsv3 {bit} -- Reserved bit for future usage. Do not use. (default: {0}) """ if not opcode in WSSettings.OPCODES: raise ValueError('Unknown opcode key') if opcode >= 0x8: # Control frames mask = 0x1 fin = 0x1 if opcode == 0x1: logging.websocket('Before encode:', data) else: logging.websocket('Before encode:', repr(data)) if opcode == 0x1: try: data = data.encode('UTF-8') except UnicodeError: raise ValueError('Text datas MUST be UTF-8 encoded.') if mask != 0x0 and mask != 0x1: raise ValueError('MASK bit parameter must be 0 or 1') if fin != 0x0 and fin != 0x1: raise ValueError('FIN bit parameter must be 0 or 1') if rsv1 != 0x0 and rsv1 != 0x1: raise ValueError('RSV1 bit parameter must be 0 or 1') if rsv2 != 0x0 and rsv2 != 0x1: raise ValueError('RSV2 bit parameter must be 0 or 1') if rsv3 != 0x0 and rsv3 != 0x1: raise ValueError('RSV3 bit parameter must be 0 or 1') if 0x3 <= opcode <= 0x7 or 0xB <= opcode: raise ValueError('Reserved opcode') bytes = struct.pack('!B', ((fin << 7) | (rsv1 << 6) | (rsv2 << 5) | (rsv3 << 4) | opcode)) mask_key = b'' if mask: # Build a random mask key (4 bytes string) for i in range(4): mask_key += chr(int(math.floor(random.random() * 256))).encode('UTF-8') logging.websocket('Mask_key:', repr(mask_key)) length = len(data) if length == 0: raise ValueError('No data given.') if length < 126: bytes += chr((mask << 7) | length).encode('UTF-8') elif length < (1 << 16): # 65536 bytes += chr((mask << 7) | 0x7e).encode('UTF-8') + struct.pack( '!H', length) elif length < (1 << 63): # 9223372036854775808 bytes += chr((mask << 7) | 0x7f).encode('UTF-8') + struct.pack( '!Q', length) else: raise ValueError('Frame too large') logging.websocket('Data=', data) ctrl = (fin, rsv1, rsv2, rsv3, opcode, mask_key, data) if mask: bytes += mask_key bytes += self.mask(mask_key, data) else: bytes += data logging.websocket('After encode:', repr(bytes)) return bytes