Add deny_new_connections and allow_new_connections

Author: Pithikos

README.md

@@ -81,9 +81,10 @@ The WebsocketServer can be initialized with the below parameters.
 | `send_message_to_all()`     | Sends a `message` to **all** connected clients. The message is a simple string.       | message         | None  |
 | `disconnect_clients_gracefully()` | Disconnect all connected clients by sending a websocket CLOSE handshake.        | Optional: status, reason | None  |
 | `disconnect_clients_abruptly()`   | Disconnect all connected clients. Clients won't be aware until they try to send some data. | None | None  |
-| `shutdown_gracefully()`     | Disconnect clients with a CLOSE handshake and shutdown server.      | Optional: status, reason | None  |
-| `shutdown_abruptly()`       | Disconnect clients and shutdown server with no handshake.      | None            | None  |
-
+| `shutdown_gracefully()`     | Disconnect clients with a CLOSE handshake and shutdown server. | Optional: status, reason      | None  |
+| `shutdown_abruptly()`       | Disconnect clients and shutdown server with no handshake.      | None                          | None  |
+| `deny_new_connections()`    | Close connection for new clients.                              | Optional: status, reason      | None  |
+| `allow_new_connections()`   | Allows back connection for new clients.                        |                               | None  |
 
 
 ### Callback functions

tests/test_server.py

@@ -141,3 +141,19 @@ def test_disconnect_clients_gracefully(self, session):
             for i in range(3):
                 client.send("test")
                 sleep(0.2)
+
+    def test_deny_new_connections(self, threaded_server):
+        url = "ws://{}:{}".format(*threaded_server.server_address)
+        server = threaded_server
+        server.deny_new_connections(status=1013, reason=b"Please try re-connecting later")
+
+        conn = websocket.create_connection(url)
+        try:
+            conn.send("test")
+        except websocket.WebSocketProtocolException as e:
+            assert 'Invalid close opcode' in e.args[0]
+        assert not server.clients
+
+        server.allow_new_connections()
+        conn = websocket.create_connection(url)
+        conn.send("test")

websocket_server/websocket_server.py

@@ -80,6 +80,12 @@ def send_message(self, 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)
 
@@ -131,6 +137,8 @@ def __init__(self, host='127.0.0.1', port=0, loglevel=logging.WARNING, key=None,
         self.id_counter = 0
         self.thread = None
 
+        self._deny_clients = False
+
     def _run_forever(self, threaded):
         cls_name = self.__class__.__name__
         try:
@@ -161,6 +169,13 @@ 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,
@@ -188,14 +203,17 @@ def handler_to_client(self, handler):
             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:
-            client["handler"].keep_alive = False
-            client["handler"].finish()
-            client["handler"].connection.close()
+            self._terminate_client_handler(client["handler"])
 
     def _shutdown_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
         """
@@ -229,6 +247,15 @@ def _disconnect_clients_abruptly(self):
         """
         self._terminate_client_handlers()
 
+    def _deny_new_connections(self, status, reason):
+        self._deny_clients = {
+            "status": status,
+            "reason": reason,
+        }
+
+    def _allow_new_connections(self):
+        self._deny_clients = False
+
 
 class WebSocketHandler(StreamRequestHandler):