# Author: Johan Hanssen Seferidis
# License: MIT
import sys
import struct
import ssl
from base64 import b64encode
from hashlib import sha1
import logging
from socket import error as SocketError
import errno
import threading
from socketserver import ThreadingMixIn, TCPServer, StreamRequestHandler
from websocket_server.thread import WebsocketServerThread
logger = logging.getLogger(__name__)
logging.basicConfig()
'''
+-+-+-+-+-------+-+-------------+-------------------------------+
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 |
+ - - - - - - - - - - - - - - - +-------------------------------+
| Payload Data continued ... |
+---------------------------------------------------------------+
'''
FIN = 0x80
OPCODE = 0x0f
MASKED = 0x80
PAYLOAD_LEN = 0x7f
PAYLOAD_LEN_EXT16 = 0x7e
PAYLOAD_LEN_EXT64 = 0x7f
OPCODE_CONTINUATION = 0x0
OPCODE_TEXT = 0x1
OPCODE_BINARY = 0x2
OPCODE_CLOSE_CONN = 0x8
OPCODE_PING = 0x9
OPCODE_PONG = 0xA
CLOSE_STATUS_NORMAL = 1000
DEFAULT_CLOSE_REASON = bytes('', encoding='utf-8')
class API():
def run_forever(self, threaded=False):
return self._run_forever(threaded)
def new_client(self, client, server):
pass
def client_left(self, client, server):
pass
def message_received(self, client, server, message):
pass
def set_fn_new_client(self, fn):
self.new_client = fn
def set_fn_client_left(self, fn):
self.client_left = fn
def set_fn_message_received(self, fn):
self.message_received = fn
def set_fn_message_received_bin(self, fn):
self.message_received_bin = fn
def send_message(self, client, msg):
self._unicast(client, msg)
def send_message_to_all(self, msg):
self._multicast(msg)
def deny_new_connections(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
self._deny_new_connections(status, reason)
def allow_new_connections(self):
self._allow_new_connections()
def shutdown_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
self._shutdown_gracefully(status, reason)
def shutdown_abruptly(self):
self._shutdown_abruptly()
def disconnect_clients_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
self._disconnect_clients_gracefully(status, reason)
def disconnect_clients_abruptly(self):
self._disconnect_clients_abruptly()
class WebsocketServer(ThreadingMixIn, TCPServer, API):
"""
A websocket server waiting for clients to connect.
Args:
port(int): Port to bind to
host(str): Hostname or IP to listen for connections. By default 127.0.0.1
is being used. To accept connections from any client, you should use
0.0.0.0.
loglevel: Logging level from logging module to use for logging. By default
warnings and errors are being logged.
Properties:
clients(list): A list of connected clients. A client is a dictionary
like below.
{
'id' : id,
'handler' : handler,
'address' : (addr, port)
}
"""
allow_reuse_address = True
daemon_threads = True # comment to keep threads alive until finished
def __init__(self, host='127.0.0.1', port=0, loglevel=logging.WARNING, key=None, cert=None):
logger.setLevel(loglevel)
TCPServer.__init__(self, (host, port), WebSocketHandler)
self.host = host
self.port = self.socket.getsockname()[1]
self.key = key
self.cert = cert
self.clients = []
self.id_counter = 0
self.thread = None
self._deny_clients = False
def _run_forever(self, threaded):
cls_name = self.__class__.__name__
try:
logger.info("Listening on port %d for clients.." % self.port)
if threaded:
self.daemon = True
self.thread = WebsocketServerThread(target=super().serve_forever, daemon=True, logger=logger)
logger.info(f"Starting {cls_name} on thread {self.thread.getName()}.")
self.thread.start()
else:
self.thread = threading.current_thread()
logger.info(f"Starting {cls_name} on main thread.")
super().serve_forever()
except KeyboardInterrupt:
self.server_close()
logger.info("Server terminated.")
except Exception as e:
logger.error(str(e), exc_info=True)
sys.exit(1)
def _message_received_(self, handler, msg):
self.message_received(self.handler_to_client(handler), self, msg)
def _message_received_bin_(self, handler, bin):
self.message_received_bin(self.handler_to_client(handler), self, bin)
def _ping_received_(self, handler, msg):
handler.send_pong(msg)
def _pong_received_(self, handler, msg):
pass
def _new_client_(self, handler):
if self._deny_clients:
status = self._deny_clients["status"]
reason = self._deny_clients["reason"]
handler.send_close(status, reason)
self._terminate_client_handler(handler)
return
self.id_counter += 1
client = {
'id': self.id_counter,
'handler': handler,
'address': handler.client_address
}
self.clients.append(client)
self.new_client(client, self)
def _client_left_(self, handler):
client = self.handler_to_client(handler)
self.client_left(client, self)
if client in self.clients:
self.clients.remove(client)
def _unicast(self, receiver_client, msg):
receiver_client['handler'].send_message(msg)
def _multicast(self, msg):
for client in self.clients:
self._unicast(client, msg)
def handler_to_client(self, handler):
for client in self.clients:
if client['handler'] == handler:
return client
def _terminate_client_handler(self, handler):
handler.keep_alive = False
handler.finish()
handler.connection.close()
def _terminate_client_handlers(self):
"""
Ensures request handler for each client is terminated correctly
"""
for client in self.clients:
self._terminate_client_handler(client["handler"])
def _shutdown_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
"""
Send a CLOSE handshake to all connected clients before terminating server
"""
self.keep_alive = False
self._disconnect_clients_gracefully(status, reason)
self.server_close()
self.shutdown()
def _shutdown_abruptly(self):
"""
Terminate server without sending a CLOSE handshake
"""
self.keep_alive = False
self._disconnect_clients_abruptly()
self.server_close()
self.shutdown()
def _disconnect_clients_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
"""
Terminate clients gracefully without shutting down the server
"""
for client in self.clients:
client["handler"].send_close(status, reason)
self._terminate_client_handlers()
