Updated script that can be controled by Nodejs web app

lib/python3.13/site-packages/selenium/webdriver/remote/websocket_connection.py

@@ -0,0 +1,146 @@
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import json
+import logging
+from ssl import CERT_NONE
+from threading import Thread
+from time import sleep
+
+from websocket import WebSocketApp  # type: ignore
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketConnection:
+    _response_wait_timeout = 30
+    _response_wait_interval = 0.1
+
+    _max_log_message_size = 9999
+
+    def __init__(self, url):
+        self.callbacks = {}
+        self.session_id = None
+        self.url = url
+
+        self._id = 0
+        self._messages = {}
+        self._started = False
+
+        self._start_ws()
+        self._wait_until(lambda: self._started)
+
+    def close(self):
+        self._ws_thread.join(timeout=self._response_wait_timeout)
+        self._ws.close()
+        self._started = False
+        self._ws = None
+
+    def execute(self, command):
+        self._id += 1
+        payload = self._serialize_command(command)
+        payload["id"] = self._id
+        if self.session_id:
+            payload["sessionId"] = self.session_id
+
+        data = json.dumps(payload)
+        logger.debug(f"-> {data}"[: self._max_log_message_size])
+        self._ws.send(data)
+
+        self._wait_until(lambda: self._id in self._messages)
+        response = self._messages.pop(self._id)
+
+        if "error" in response:
+            raise Exception(response["error"])
+        else:
+            result = response["result"]
+            return self._deserialize_result(result, command)
+
+    def add_callback(self, event, callback):
+        event_name = event.event_class
+        if event_name not in self.callbacks:
+            self.callbacks[event_name] = []
+
+        def _callback(params):
+            callback(event.from_json(params))
+
+        self.callbacks[event_name].append(_callback)
+        return id(_callback)
+
+    on = add_callback
+
+    def remove_callback(self, event, callback_id):
+        event_name = event.event_class
+        if event_name in self.callbacks:
+            for callback in self.callbacks[event_name]:
+                if id(callback) == callback_id:
+                    self.callbacks[event_name].remove(callback)
+                    return
+
+    def _serialize_command(self, command):
+        return next(command)
+
+    def _deserialize_result(self, result, command):
+        try:
+            _ = command.send(result)
+            raise Exception("The command's generator function did not exit when expected!")
+        except StopIteration as exit:
+            return exit.value
+
+    def _start_ws(self):
+        def on_open(ws):
+            self._started = True
+
+        def on_message(ws, message):
+            self._process_message(message)
+
+        def on_error(ws, error):
+            logger.debug(f"error: {error}")
+            ws.close()
+
+        def run_socket():
+            if self.url.startswith("wss://"):
+                self._ws.run_forever(sslopt={"cert_reqs": CERT_NONE}, suppress_origin=True)
+            else:
+                self._ws.run_forever(suppress_origin=True)
+
+        self._ws = WebSocketApp(self.url, on_open=on_open, on_message=on_message, on_error=on_error)
+        self._ws_thread = Thread(target=run_socket)
+        self._ws_thread.start()
+
+    def _process_message(self, message):
+        message = json.loads(message)
+        logger.debug(f"<- {message}"[: self._max_log_message_size])
+
+        if "id" in message:
+            self._messages[message["id"]] = message
+
+        if "method" in message:
+            params = message["params"]
+            for callback in self.callbacks.get(message["method"], []):
+                callback(params)
+
+    def _wait_until(self, condition):
+        timeout = self._response_wait_timeout
+        interval = self._response_wait_interval
+
+        while timeout > 0:
+            result = condition()
+            if result:
+                return result
+            else:
+                timeout -= interval
+                sleep(interval)