diff --git a/bec_widgets/cli/client_utils.py b/bec_widgets/cli/client_utils.py index 06050c15..91273cf8 100644 --- a/bec_widgets/cli/client_utils.py +++ b/bec_widgets/cli/client_utils.py @@ -11,9 +11,11 @@ from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING +import msgpack from bec_lib.endpoints import MessageEndpoints from bec_lib.logger import bec_logger from bec_lib.utils.import_utils import isinstance_based_on_class_name, lazy_import, lazy_import_from +from cryptography.fernet import Fernet import bec_widgets.cli.client as client from bec_widgets.cli.auto_updates import AutoUpdates @@ -67,7 +69,14 @@ def _get_output(process, logger) -> None: logger.error(f"Error reading process output: {str(e)}") -def _start_plot_process(gui_id: str, gui_class: type, config: dict | str, logger=None) -> None: +def _start_plot_process( + gui_id: str, + gui_class: type, + config: dict | str, + acl_data: bytes | None = None, + token: bytes | None = None, + logger=None, +) -> None: """ Start the plot in a new process. @@ -84,6 +93,8 @@ def _start_plot_process(gui_id: str, gui_class: type, config: dict | str, logger env_dict = os.environ.copy() env_dict["PYTHONUNBUFFERED"] = "1" + env_dict["BEC_GUI_ACL"] = acl_data or b"" + env_dict["BEC_GUI_TOKEN"] = token or b"" if logger is None: stdout_redirect = subprocess.DEVNULL @@ -179,6 +190,7 @@ class BECGuiClient(RPCBase): self._gui_started_event = threading.Event() self._process = None self._process_output_processing_thread = None + self._fernet = None @property def windows(self): @@ -263,6 +275,18 @@ class BECGuiClient(RPCBase): self._do_show_all() self._gui_started_event.set() + # def _update_gui_acls(self, username: str, password: str | None): + # self._client.connector.send( + # MessageEndpoints.gui_acls(self._gui_id), + # messages.CredentialsMessage( + # credentials={ + # "token": self._fernet.encrypt( + # msgpack.dumps({"username": username, "password": password}) + # ) + # } + # ), + # ) + def start_server(self, wait=False) -> None: """ Start the GUI server, and execute callback when it is launched @@ -271,8 +295,18 @@ class BECGuiClient(RPCBase): logger.success("GUI starting...") self._startup_timeout = 5 self._gui_started_event.clear() + encr_token = Fernet.generate_key() + self._fernet = Fernet(encr_token) + conn = self._client.connector._redis_conn.connection_pool.connection_kwargs + acl_data = {"username": conn.get("username"), "password": conn.get("password")} + acl_data = self._fernet.encrypt(msgpack.dumps(acl_data)) self._process, self._process_output_processing_thread = _start_plot_process( - self._gui_id, self.__class__, self._client._service_config.config, logger=logger + self._gui_id, + self.__class__, + self._client._service_config.config, + acl_data=acl_data, + token=encr_token, + logger=logger, ) def gui_started_callback(callback): diff --git a/bec_widgets/cli/server.py b/bec_widgets/cli/server.py index 1f854f25..ff403757 100644 --- a/bec_widgets/cli/server.py +++ b/bec_widgets/cli/server.py @@ -2,16 +2,19 @@ from __future__ import annotations import functools import json +import os import signal import sys import types from contextlib import contextmanager, redirect_stderr, redirect_stdout from typing import Union +import msgpack from bec_lib.endpoints import MessageEndpoints from bec_lib.logger import bec_logger from bec_lib.service_config import ServiceConfig from bec_lib.utils.import_utils import lazy_import +from cryptography.fernet import Fernet from qtpy.QtCore import Qt, QTimer from redis.exceptions import RedisError @@ -57,7 +60,10 @@ class BECWidgetsCLIServer: client=None, config=None, gui_class: Union[BECFigure, BECDockArea] = BECFigure, + token: str = None, ) -> None: + self._fernet = Fernet(token) if token else None + self._init_acls(config) self.status = messages.BECStatus.BUSY self.dispatcher = BECDispatcher(config=config) if dispatcher is None else dispatcher self.client = self.dispatcher.client if client is None else client @@ -67,6 +73,8 @@ class BECWidgetsCLIServer: self.rpc_register = RPCRegister() self.rpc_register.add_rpc(self.gui) + # self.dispatcher.connect_slot(self.on_acl_update, MessageEndpoints.gui_acl(self.gui_id)) + self.dispatcher.connect_slot( self.on_rpc_update, MessageEndpoints.gui_instructions(self.gui_id) ) @@ -79,6 +87,13 @@ class BECWidgetsCLIServer: self.status = messages.BECStatus.RUNNING logger.success(f"Server started with gui_id: {self.gui_id}") + def _init_acls(self, config: ServiceConfig): + acl_data = os.getenv("BEC_GUI_ACL") + if not acl_data: + return + acl_data = msgpack.loads(self._fernet.decrypt(acl_data)) + config.config["acl"] = acl_data + def on_rpc_update(self, msg: dict, metadata: dict): request_id = metadata.get("request_id") logger.debug(f"Received RPC instruction: {msg}, metadata: {metadata}") @@ -96,6 +111,9 @@ class BECWidgetsCLIServer: logger.debug(f"RPC instruction executed successfully: {res}") self.send_response(request_id, True, {"result": res}) + def on_acl_update(self, msg: dict, metadata: dict): + logger.debug(f"Received ACL update: {msg}, metadata: {metadata}") + def send_response(self, request_id: str, accepted: bool, msg: dict): self.client.connector.set_and_publish( MessageEndpoints.gui_instruction_response(request_id), @@ -190,13 +208,10 @@ def _start_server(gui_id: str, gui_class: Union[BECFigure, BECDockArea], config: # if no config is provided, use the default config service_config = ServiceConfig() - # bec_logger.configure( - # service_config.redis, - # QtRedisConnector, - # service_name="BECWidgetsCLIServer", - # service_config=service_config.service_config, - # ) - server = BECWidgetsCLIServer(gui_id=gui_id, config=service_config, gui_class=gui_class) + token = os.getenv("BEC_GUI_TOKEN") + server = BECWidgetsCLIServer( + gui_id=gui_id, config=service_config, gui_class=gui_class, token=token + ) return server diff --git a/pyproject.toml b/pyproject.toml index b4539cb8..0691f4a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "pyte", # needed for vt100 console "qtconsole~=5.5, >=5.5.1", # needed for jupyter console "qtpy~=2.4", + "cryptography~=44.0", ]