Private model webui (#666)

* add decryption file and tooltips/placeholders to container registration

* change profile -> settings page, refactor according to changes

* add certificate settings

* modify GetUserCertificate.run - interactive

* add certificate settings in webui

* fix certificate corruption

* tmp

* 'Get' instead of 'Generate' certificate - certificate status 'local' to 'to be uploaded'

* fix database lock error after token expiration in default profile

* join thread after setting the event - auto grant access

* fix certificate corruption - updated

* modify 'generate' into get/retrieve certificate - settings.py. submit certificate without passing name

* use sanitize path for decryption key and add more checks for validation

* tmp

* make 'access pending' underlined in dataset details - disable model run button if access isn't granted

* - add emails to grant access and auto grant access.
- show keys with revoke btn in container_access page, add revoke functionality.

* move get_container_type and generate_uuid to webui.utils

* fix decryption_key error in submit.py

* add regex in UI to filter url/code messages from get_client_certificate output

* refactor notifications/add critical printing (popup) in webui

* refactor modals - making print critical works. add modal queue

* fix some ui bugs - contianer_access.html

* update e2e tests according to modal changes
add:
- encrypted container registration/association
- generate and submit certificate for dataset owner
- grant access - encrypted model owner to dataset owner
- delete keys after running the flow

* fix critical print

* make settings page accessible if user is not logged in.
merge CA and Fingerprint settings into profile settings (platform, gpus)
modify some ui bugs.

* fix e2e tests

* fix e2e tests - profile activation

* fix certificate checks - unit tests

* rename ui.regex - ui.is_parsed_output

* add decryption key to compatibility test form - webui

* refactor parsed_ouput

* fix cert issues

* fix codeql alerts for decryption key

* fix set profile args logic in webui

* fix redirect security alert

* fix codeql error

* fix e2e test - compatibility test

---------

Co-authored-by: hasan7n <hasankassim7@hotmail.com>

Author: mhmdk0

cli/medperf/commands/certificate/client_certificate.py

@@ -1,6 +1,7 @@
 from medperf.account_management import get_medperf_user_data
 from medperf.certificates import get_client_cert
 from medperf.exceptions import MedperfException
+from medperf.ui.interface import UI
 from medperf.utils import get_pki_assets_path, remove_path
 from medperf import config
 from medperf.entities.ca import CA
@@ -11,6 +12,7 @@ class GetUserCertificate:
     @staticmethod
     def run(overwrite: bool = False):
         """get user cert"""
+        ui: UI = config.ui
         ca_id = config.certificate_authority_id
         ca = CA.get(ca_id)
 
@@ -22,4 +24,8 @@ def run(overwrite: bool = False):
                     "Cert and key already present. Rerun the command with --overwrite"
                 )
             remove_path(output_path, sensitive=True)
-        get_client_cert(ca, email, output_path)
+
+        with ui.interactive():
+            ui.text = "Getting Certificate"
+            with ui.parsed_output():
+                get_client_cert(ca, email, output_path)

cli/medperf/commands/certificate/utils.py

@@ -1,12 +1,25 @@
 from medperf import config
 from medperf.entities.certificate import Certificate
-from medperf.utils import get_pki_assets_path
+from medperf.utils import get_pki_assets_path, remove_path
 from medperf.account_management import get_medperf_user_data
 import os
 import base64
 import logging
 
 
+def _check_and_clean_certificate_corruption(local_cert_folder):
+    private_key_path = os.path.join(local_cert_folder, config.private_key_file)
+    certificate_path = os.path.join(local_cert_folder, config.certificate_file)
+
+    private_key_exists = os.path.exists(private_key_path)
+    certificate_exists = os.path.exists(certificate_path)
+
+    certificate_corrupted = not private_key_exists or not certificate_exists
+
+    if certificate_corrupted:
+        remove_path(local_cert_folder, sensitive=True)
+
+
 def current_user_certificate_status():
     """Check the status of the current user certificate. Possible cases:
     - No local certificate folder and no submitted certificate
@@ -35,6 +48,9 @@ def current_user_certificate_status():
     email = get_medperf_user_data()["email"]
     local_cert_folder = get_pki_assets_path(email, config.certificate_authority_id)
 
+    # If certificate is corrupted, delete the certificate folder
+    _check_and_clean_certificate_corruption(local_cert_folder)
+
     # Check
     exists_locally = os.path.exists(local_cert_folder)
     submitted = user_cert_object is not None

cli/medperf/commands/compatibility_test/run.py

@@ -6,6 +6,7 @@
 from medperf.entities.benchmark import Benchmark
 from medperf.entities.report import TestReport
 from medperf.exceptions import InvalidArgumentError
+from medperf.utils import sanitize_path
 from .validate_params import CompatibilityTestParamsValidator
 from .utils import download_demo_data, prepare_cube, get_cube, create_test_dataset
 import medperf.config as config
@@ -29,7 +30,7 @@ def run(
         skip_data_preparation_step: bool = False,
         use_local_model_image: bool = False,
         model_decryption_key: Path = None,
-    ) -> (str, dict):
+    ) -> tuple[str, dict]:
         """Execute a test workflow. Components of a complete workflow should be passed.
         When only the benchmark is provided, it implies the following workflow will be used:
         - the benchmark's demo dataset is used as the raw data
@@ -143,7 +144,7 @@ def __init__(
         self.evaluator_cube = None
 
         # Decryption key is used for compatibility test of encrypted containers
-        self.model_decryption_key = model_decryption_key
+        self.model_decryption_key = sanitize_path(model_decryption_key)
 
         self.validator = CompatibilityTestParamsValidator(
             self.benchmark_uid,

cli/medperf/commands/mlcube/grant_access.py

@@ -71,11 +71,12 @@ def validate_allowed_emails(self):
             self.allowed_emails = validate_and_normalize_emails(self.allowed_emails)
 
     def verify_certificate_authority(self):
-        config.ui.print("Verifying Certificate Authority")
         ca = CA.get(uid=config.certificate_authority_id)
-        verify_certificate_authority(
-            ca, expected_fingerprint=config.certificate_authority_fingerprint
-        )
+        with config.ui.interactive():
+            config.ui.text = "Verifying Certificate Authority"
+            verify_certificate_authority(
+                ca, expected_fingerprint=config.certificate_authority_fingerprint
+            )
 
     def prepare_certificates_list(self):
         config.ui.print("Getting Data Owner Certificates")

cli/medperf/commands/mlcube/submit.py

@@ -3,7 +3,7 @@
 import medperf.config as config
 from medperf.entities.cube import Cube
 from medperf.exceptions import InvalidArgumentError
-from medperf.utils import remove_path, store_decryption_key
+from medperf.utils import remove_path, sanitize_path, store_decryption_key
 import logging
 
 
@@ -33,7 +33,7 @@ def __init__(self, submit_info: dict, decryption_key: str = None):
         self.comms = config.comms
         self.ui = config.ui
         self.cube = Cube(**submit_info)
-        self.decryption_key = decryption_key
+        self.decryption_key = sanitize_path(decryption_key)
         config.tmp_paths.append(self.cube.path)
 
     def download_config_files(self):
@@ -50,6 +50,16 @@ def validate(self):
                 "Container is not encrypted, but a decryption key is provided"
             )
 
+        if self.decryption_key is not None:
+            if not os.path.exists(self.decryption_key):
+                raise InvalidArgumentError(
+                    f"Decryption key does not exist at path {self.decryption_key}"
+                )
+            if os.path.isdir(self.decryption_key):
+                raise InvalidArgumentError(
+                    f"The provided decryption key path {self.decryption_key} is a directory!"
+                )
+
     def download_run_files(self):
         self.cube.download_run_files()
 

cli/medperf/comms/auth/auth0.py

@@ -174,11 +174,15 @@ def access_token(self):
             except sqlite3.OperationalError:
                 msg = "Another process is using the database. Try again later"
                 raise CommunicationError(msg)
-            token = self._access_token
+
             # Sqlite will automatically execute COMMIT and close the connection
             # if an exception is raised during the retrieval of the access token.
-            db.execute("COMMIT")
-            db.close()
+
+            try:
+                token = self._access_token
+            finally:
+                db.execute("COMMIT")
+                db.close()
 
             return token
 

cli/medperf/tests/commands/certificate/test_utils_certificates.py

@@ -99,11 +99,21 @@ def test_current_user_certificate_status(
     mocker.patch(PATCH_STATUS.format("get_pki_assets_path"), return_value=local_path)
     if local_exists:
         fs.create_dir(local_path)
+        fs.create_file(
+            os.path.join(local_path, config.private_key_file), contents=b"private-key"
+        )
+        fs.create_file(
+            os.path.join(local_path, config.certificate_file), contents=b"abc"
+        )
         if remote_exists and local_matches:
+            if os.path.exists(os.path.join(local_path, config.certificate_file)):
+                fs.remove(os.path.join(local_path, config.certificate_file))
             fs.create_file(
                 os.path.join(local_path, config.certificate_file), contents=b"abc"
             )
         elif remote_exists and not local_matches:
+            if os.path.exists(os.path.join(local_path, config.certificate_file)):
+                fs.remove(os.path.join(local_path, config.certificate_file))
             fs.create_file(
                 os.path.join(local_path, config.certificate_file), contents=b"xyz"
             )

cli/medperf/ui/cli.py

@@ -1,3 +1,4 @@
+import re
 import typer
 from getpass import getpass
 from yaspin import yaspin
@@ -10,13 +11,56 @@ class CLI(UI):
     def __init__(self):
         self.spinner = yaspin(color="green")
         self.is_interactive = False
+        self.is_parsed_output = False
+
+    def print_url_message(self, message: str):
+        match = re.search(r"https?://[^\s]+", message)
+
+        url = match.group(0)
+        start, end = match.span()
+
+        before = message[:start].strip()
+        after = message[end:].strip()
+
+        if before.strip():
+            self.print(before)
+
+        self.print_url(url)
+
+        if after.strip():
+            self.print(after)
+
+    def contains_url(self, message: str):
+        return bool(re.search(r"https?://[^\s]+", message))
+
+    def is_code(self, message: str):
+        ansi_escape = re.compile(r"\x1b\[[0-9;]*m")
+        message = ansi_escape.sub("", message.strip())
+        return bool(re.fullmatch(r"[A-Z0-9]{4}-[A-Z0-9]{4}-[A-Z0-9]{4}", message))
+
+    def print_subprocess_logs(self, msg: str):
+        """Display subprocess logs on the command line
+
+        Args:
+            msg (str): message to print
+        """
+        if self.is_parsed_output:
+            if self.is_code(msg):
+                self.print_code(msg)
+            elif self.contains_url(msg):
+                self.print_url_message(msg)
+            else:
+                self.print(msg)
+        else:
+            self.print(msg)
 
     def print(self, msg: str = ""):
         """Display a message on the command line
 
         Args:
             msg (str): message to print
         """
+
         self._print(msg)
 
     def print_error(self, msg: str):
@@ -29,6 +73,9 @@ def print_error(self, msg: str):
         msg = typer.style(msg, fg=typer.colors.RED, bold=True)
         self._print(msg)
 
+    def print_critical(self, msg: str):
+        self.print_warning(msg)
+
     def print_warning(self, msg: str):
         """Display a warning message on the command line
 
@@ -78,6 +125,31 @@ def interactive(self):
             finally:
                 self.stop_interactive()
 
+    def start_parsed_output(self):
+        """Start a parsed output session where messages will be displayed based on regular expressions"""
+        self.is_parsed_output = True
+
+    def stop_parsed_output(self):
+        """Stop the parsed output session"""
+        self.is_parsed_output = False
+
+    @contextmanager
+    def parsed_output(self):
+        """Context managed parsed output session.
+
+        Yields:
+            CLI: Yields the current CLI instance with a parsed output session initialized
+        """
+        if self.is_parsed_output:
+            # if already parsed output, do nothing
+            yield self
+        else:
+            self.start_parsed_output()
+            try:
+                yield self
+            finally:
+                self.stop_parsed_output()
+
     @property
     def text(self):
         return self.spinner.text