aixtools 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.8'
-__version_tuple__ = version_tuple = (0, 1, 8)
+__version__ = version = '0.1.9'
+__version_tuple__ = version_tuple = (0, 1, 9)
 
 __commit_id__ = commit_id = None

aixtools/compliance/private_data.py

@@ -0,0 +1,136 @@
+import json
+from pathlib import Path
+
+from fastmcp import Context
+
+from aixtools.server.path import get_workspace_path
+
+PRIVATE_DATA_FILE = ".private_data"
+
+
+class PrivateData:
+    """
+    Class to manage private data file in the workspace.
+
+    The information is stored in a JSON file named `.private_data` within the workspace directory.
+    If the file does not exist, it indicates that there is no private data.
+
+    IMPORTANT: All modifications save the data to the file immediately.
+
+    FIXME: We should add some level of mutex/locking to prevent concurrent writes.
+    """
+
+    def __init__(self, ctx: Context | None = None):
+        self.ctx: Context | None = ctx
+        self._has_private_data: bool = False  # Flag indicating if private data exists
+        self._private_datasets: list[str] = []  # List of private datasets
+        self._idap_datasets: list[str] = []  # List of dataset with IDAP
+        self.load()
+
+    def add_private_dataset(self, dataset_name: str) -> None:
+        """
+        Add a private dataset to the list.
+        Save the state after modification.
+        """
+        if dataset_name not in self._private_datasets:
+            self._private_datasets.append(dataset_name)
+            self._has_private_data = True
+            self.save()
+
+    def add_idap_dataset(self, dataset_name: str) -> None:
+        """
+        Add a dataset with IDAP to the list.
+        This also adds it to the private datasets if not already present.
+        Save the state after modification.
+        """
+        if not self.has_idap_dataset(dataset_name):
+            self._idap_datasets.append(dataset_name)
+            self._has_private_data = True
+            # An IDAP dataset is also a private dataset
+            if not self.has_private_dataset(dataset_name):
+                self._private_datasets.append(dataset_name)
+            self.save()
+
+    def get_private_datasets(self) -> list[str]:
+        """Get the list of private datasets as a copy (to avoid modification)."""
+        return list(self._private_datasets)
+
+    def get_idap_datasets(self) -> list[str]:
+        """Get the list of datasets with IDAP as a copy (to avoid modification)."""
+        return list(self._idap_datasets)
+
+    def has_private_dataset(self, dataset_name: str) -> bool:
+        """Check if a specific private dataset exists."""
+        return dataset_name in self._private_datasets
+
+    def has_idap_dataset(self, dataset_name: str) -> bool:
+        """Check if a specific dataset with IDAP exists."""
+        return dataset_name in self._idap_datasets
+
+    @property
+    def has_private_data(self) -> bool:
+        """Check if private data exists."""
+        return self._has_private_data
+
+    @has_private_data.setter
+    def has_private_data(self, value: bool) -> None:
+        """
+        Set the flag indicating if private data exists.
+        Save the state after modification.
+        """
+        self._has_private_data = value
+        if not value:
+            self._private_datasets = []
+            self._idap_datasets = []
+        self.save()
+
+    def _get_private_data_path(self) -> Path:
+        """Get the path to the private data file in the workspace."""
+        return get_workspace_path(service_name=None, ctx=self.ctx) / PRIVATE_DATA_FILE
+
+    def _has_private_data_file(self) -> bool:
+        """Check if the private data file exists in the workspace."""
+        private_data_path = self.get_private_data_path()  # type: ignore
+        return private_data_path.exists()
+
+    def save(self) -> None:
+        """Save content to the private data file in the workspace."""
+        private_data_path = self._get_private_data_path()
+        # No private data? Delete the file if it exists
+        if not self.has_private_data:
+            private_data_path.unlink(missing_ok=True)
+            return
+        # If there is private data, serialize this object as JSON
+        private_data_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(private_data_path, "w") as f:
+            # Dump class as JSON, excluding the context
+            data_dict = self.__dict__.copy()
+            data_dict["ctx"] = None
+            json_data = json.dumps(data_dict, indent=4)
+            f.write(json_data)
+
+    def load(self) -> None:
+        """Load content from the private data file in the workspace."""
+        private_data_path = self._get_private_data_path()
+        if not private_data_path.exists():
+            # No private data file
+            self.has_private_data = False
+            self._private_datasets = []
+            self._idap_datasets = []
+            return
+        with open(private_data_path, "r") as f:
+            data = json.load(f)
+            self.has_private_data = data.get("_has_private_data", False)
+            self._private_datasets = data.get("_private_datasets", [])
+            self._idap_datasets = data.get("_idap_datasets", [])
+
+    def __repr__(self) -> str:
+        return (
+            f"PrivateData(has_private_data={self.has_private_data}, "
+            f"private_datasets={self._private_datasets}, "
+            f"idap_datasets={self._idap_datasets}), "
+            f"file_path={self._get_private_data_path()})"
+        )
+
+    def __str__(self) -> str:
+        return self.__repr__()

aixtools/server/__init__.py

@@ -13,6 +13,10 @@ from .utils import (
     get_session_id_tuple,
     run_in_thread,
 )
+from .workspace_privacy import (
+    is_session_private,
+    set_session_private,
+)
 
 __all__ = [
     "get_workspace_path",
@@ -20,4 +24,6 @@ __all__ = [
     "container_to_host_path",
     "host_to_container_path",
     "run_in_thread",
+    "is_session_private",
+    "set_session_private",
 ]

aixtools/server/workspace_privacy.py

@@ -0,0 +1,65 @@
+"""
+Workspace privacy utilities for managing session-level privacy flags.
+"""
+
+from pathlib import Path
+
+from fastmcp import Context
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.server.path import get_workspace_path
+
+logger = get_logger(__name__)
+
+PRIVACY_FLAG_FILENAME = ".private_data_indicator"
+
+
+def set_session_private(ctx: Context | tuple[str, str] | None = None) -> bool:
+    """
+    Set the current session as private by creating a privacy flag file.
+
+    Creates an empty file in the session workspace directory
+    and sets it as read-only to prevent accidental removal
+
+    Args:
+        ctx: FastMCP context for user/session identification.
+             If None, uses current FastMCP request context from HTTP headers.
+             If tuple, first part is a user id (username), second part is session id (aka conversation id)
+
+    Returns:
+        bool: True if privacy flag was successfully created, False otherwise.
+    """
+    try:
+        workspace_path = Path(get_workspace_path(ctx=ctx))
+        privacy_file = workspace_path / PRIVACY_FLAG_FILENAME
+        workspace_path.mkdir(parents=True, exist_ok=True)
+        privacy_file.touch(exist_ok=True)
+        privacy_file.chmod(0o444)
+        logger.warning("Session marked as private")
+        return True
+    except (OSError, ValueError, RuntimeError) as e:
+        logger.error("Set current session as private: %s", str(e))
+        return False
+
+
+def is_session_private(ctx: Context | tuple[str, str] | None = None) -> bool:
+    """
+    Check if the current session is marked as private.
+
+    Args:
+        ctx: FastMCP context for user/session identification.
+             If None, uses current FastMCP request context from HTTP headers.
+             If tuple, first part is a user id (username), second part is session id (aka conversation id)
+
+    Returns:
+        bool: True if workspace is private (flag file exists), False otherwise.
+    """
+    try:
+        workspace_path = Path(get_workspace_path(ctx=ctx))
+        privacy_file = workspace_path / PRIVACY_FLAG_FILENAME
+        is_private = privacy_file.exists()
+        logger.info("Session privacy check, is private: %s", str(is_private))
+        return is_private
+    except (OSError, ValueError, RuntimeError) as e:
+        logger.error("Check session privacy: %s, assuming not private!", str(e))
+        return False

aixtools/vault/vault.py

@@ -27,9 +27,11 @@ class VaultClient:
         if not self.client.is_authenticated():
             raise VaultAuthError("Vault client authentication failed. Check vault_token.")
 
-    def _get_secret_path(self, user_id: str, service_name: str) -> str:
-        """Generate the vault secret path for a user and service."""
-        return f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+    def _get_secret_path(self, user_id: str, service_name: Optional[str] = None) -> str:
+        """Generate the vault secret path for a user and optionally a service."""
+        if service_name:
+            return f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+        return f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}"
 
     def store_user_service_api_key(self, *, user_id: str, service_name: str, user_api_key: str):
         """
@@ -98,3 +100,38 @@ class VaultClient:
         except Exception as e:
             logger.error("Failed to read complete secret from path %s: %s", secret_path, str(e))
             raise VaultAuthError(e) from e
+
+    def list_user_secret_keys(self, *, user_id: str) -> list[str]:
+        """
+        List all secret keys (service names) for a user, optionally filtered by service name.
+
+        Args:
+            user_id: The user ID to list secrets for
+            service_name: Optional service name to filter results. If provided, returns only this service if it exists.
+
+        Returns:
+            List of service names (secret keys) for the user. Empty list if no secrets exist.
+        """
+        try:
+            # List all services for user
+            user_path = self._get_secret_path(user_id)
+            logger.info("Listing secret keys for user at path %s", user_path)
+
+            response = self.client.secrets.kv.v2.list_secrets(path=user_path, mount_point=config.VAULT_MOUNT_POINT)
+
+            if response and "data" in response and "keys" in response["data"]:
+                secret_keys = response["data"]["keys"]
+                # Remove trailing slashes from directory names if any
+                secret_keys = [key.rstrip("/") for key in secret_keys]
+                logger.info("Found %d secret keys for user %s", len(secret_keys), user_id)
+                return secret_keys
+            logger.info("No secret keys found for user %s", user_id)
+            return []
+
+        except InvalidPath:
+            # User path does not exist
+            logger.warning("User path does not exist for user %s", user_id)
+            return []
+        except Exception as e:
+            logger.error("Failed to list secret keys for user %s: %s", user_id, str(e))
+            raise VaultAuthError(e) from e

{aixtools-0.1.8.dist-info → aixtools-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.1.8
+Version: 0.1.9
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.1.8.dist-info → aixtools-0.1.9.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=Zaz3s9gl_rzsS46-ymJOALojMxviW77EJq_agE8knLk,704
+aixtools/_version.py,sha256=ib8ckvf-NNDfacXd8unW0p5cf-gl57XyQvjoEMc_pvc,704
 aixtools/app.py,sha256=JzQ0nrv_bjDQokllIlGHOV0HEb-V8N6k_nGQH-TEsVU,5227
 aixtools/chainlit.md,sha256=yC37Ly57vjKyiIvK4oUvf4DYxZCwH7iocTlx7bLeGLU,761
 aixtools/context.py,sha256=I_MD40ZnvRm5WPKAKqBUAdXIf8YaurkYUUHSVVy-QvU,598
@@ -29,6 +29,7 @@ aixtools/agents/__init__.py,sha256=MAW196S2_G7uGqv-VNjvlOETRfuV44WlU1leO7SiR0A,2
 aixtools/agents/agent.py,sha256=E1zu70t53RqIbcLI_R09wUtsiYZR1bTnElCQ5PrsrKw,6127
 aixtools/agents/agent_batch.py,sha256=0Zu9yNCRPAQZPjXQ-dIUAmP1uGTVbxVt7xvnMpoJMjU,2251
 aixtools/agents/prompt.py,sha256=VCOVSnhNKsPIT347ouzwM1PH4I9UTm2cSnTh3ZpjRwk,3391
+aixtools/compliance/private_data.py,sha256=aYEd13eHUajs72nQxCXsa7tklXnqyurUaYuNzm8belQ,5258
 aixtools/db/__init__.py,sha256=b8vRhme3egV-aUZbAntnOaDkSXB8UT0Xy5oqQhU_z0Q,399
 aixtools/db/database.py,sha256=caWe95GlxZYlxn2ubDmR-_cQUW0ulkpR3BHunKIaOsw,3369
 aixtools/db/vector_db.py,sha256=be4JGyXj3o8VEfy9L6SO1aAoDET_zazMJkYfjlYHTYQ,4133
@@ -56,10 +57,11 @@ aixtools/mcp/example_server.py,sha256=1SWCyrLWsAnOa81HC4QbPJo_lBVu0b3SZBWI-qDh1v
 aixtools/mcp/fast_mcp_log.py,sha256=XYOS406dVjn5YTHyGRsRvVNQ0SKlRObfrKj6EeLFjHg,1057
 aixtools/mcp/faulty_mcp.py,sha256=uU9vlNGCS_i2k20wocVMaDHTlYjMQxuzjILad9O1cjA,12807
 aixtools/model_patch/model_patch.py,sha256=JT-oHubIn2LeoSwWbzEQ5vLH7crJmFUecHyQfaAFHa0,1813
-aixtools/server/__init__.py,sha256=rwPx020YpOzCnrxA80Lc4yLLcIp-Mpe9hNqVO9wDPv0,448
+aixtools/server/__init__.py,sha256=37ADJrGLzsmjFsM2ZKUoM9cevH8rBn359WesDxIwoco,585
 aixtools/server/app_mounter.py,sha256=0tJ0tC140ezAjnYdlhpLJQjY-TO8NVw7D8LseYCCVY8,3336
 aixtools/server/path.py,sha256=SaIJxvmhJy3kzx5zJ6d4cKP6kKu2wFFciQkOLGTA4gg,3056
 aixtools/server/utils.py,sha256=tZWITIx6M-luV9yve4j3rPtYGSSA6zWS0JWEAySne_M,2276
+aixtools/server/workspace_privacy.py,sha256=grcj82eHSd7gFbb5f_w9nv4TWp50QyU952l0iIPoChM,2375
 aixtools/testing/__init__.py,sha256=mlmaAR2gmS4SbsYNCxnIprmFpFp-syjgVUkpUszo3mE,166
 aixtools/testing/aix_test_model.py,sha256=dlI3sdyvmu4fUs_K4-oazs_a7cE6V-gnI6RQ0_fPVxg,5925
 aixtools/testing/mock_tool.py,sha256=4I0LxxSkLhGIKM2YxCP3cnYI8IYJjdKhfwGZ3dioXsM,2465
@@ -77,8 +79,9 @@ aixtools/utils/utils.py,sha256=5911Ej1ES2NU_FKIWA3CWKhKnwgjvi1aDR2aiD6Xv3E,4880
 aixtools/utils/chainlit/cl_agent_show.py,sha256=vaRuowp4BRvhxEr5hw0zHEJ7iaSF_5bo_9BH7pGPPpw,4398
 aixtools/utils/chainlit/cl_utils.py,sha256=fxaxdkcZg6uHdM8uztxdPowg3a2f7VR7B26VPY4t-3c,5738
 aixtools/vault/__init__.py,sha256=fsr_NuX3GZ9WZ7dGfe0gp_5-z3URxAfwVRXw7Xyc0dU,141
-aixtools/vault/vault.py,sha256=WkzBTEYM-Vqjyoa5x5imbEDi0ePBklMjD_aAdvIK-34,4293
-docker/mcp-base/Dockerfile,sha256=UuAA1ltJpusKwow2fd_wS0M4DexbxyCzQGNVOK1_lzs,1547
+aixtools/vault/vault.py,sha256=9dZLWdZQk9qN_Q9Djkofw9LUKnJqnrX5H0fGusVLBhA,6037
+docker/mcp-base/Dockerfile,sha256=uislVoTEgRF--AAiyX24sBxlDdfA1ZU5rDM94XYFqvI,1388
+docker/mcp-base/zscaler.crt,sha256=fCUNiOfJlWTA7R4zV1Xyb-XC1_nMHPoYTFGvBj1oz6s,1732
 notebooks/example_faulty_mcp_server.ipynb,sha256=b2Cy3GXfj-gOBZ7SoUzj25F1rxp5u-32EWPHWQ-sxn8,1729
 notebooks/example_mcp_server_stdio.ipynb,sha256=ya4dRKNFU2vQxob-uIhKHGAzINXGQ6MehgKVmSCpHLk,1634
 notebooks/example_raw_mcp_client.ipynb,sha256=uchaG-LuuwEpE2oIkmhZ2s1EDb19AgT1fUv2Jxtjgu8,1795
@@ -89,7 +92,7 @@ scripts/log_view.sh,sha256=bp8oXFRRbbHpyvHAN85wfDHTVK7vMJOYsBx_-bgECQc,511
 scripts/run_example_mcp_server.sh,sha256=f7m7h7O_wo6-nAsYlOXVWIASCOh3Qbuu0XWizlxMhl8,355
 scripts/run_faulty_mcp_server.sh,sha256=u_-8NbPDnJQt6IinNSjh8tc2ed-_MjGyipJXrUXaGR8,291
 scripts/run_server.sh,sha256=5iiB9bB5M2MuOgxVQqu7Oa_tBVtJpt0uB4z9uLu2J50,720
-scripts/test.sh,sha256=5akwfnX7P-98898WCXcQRFQt84UBgms5Y9fXC_ym1a4,662
+scripts/test.sh,sha256=KxXWkVqctFRNP8hItJr8K27nDHEkfwNWb1UFhpBQDOk,865
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/a2a/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -102,6 +105,7 @@ tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_agent_executor.py,sha256=PcyC
 tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_storage.py,sha256=tb67pFfvyWSaDfKaiPDNBQfl6-o17WtCMZh3lQHrYxY,5468
 tests/unit/agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/agents/test_prompt.py,sha256=YWFZdH_F774hxw79gsWoTWBPVs8UjOAtJOgNXJ8N9gs,15384
+tests/unit/compliance/test_private_data.py,sha256=GjH7NCp54Bz1S-CmH_mUe53lb53kllOOJEm448OniRI,13693
 tests/unit/google/__init__.py,sha256=eRYHldBi5cFWL7oo2_t5TErI8ESmIjNvBZIcp-w8hSA,45
 tests/unit/google/test_client.py,sha256=fXR4Cozea7bdL2prM-1s9IqUQ9AheklQnHpN-4YM3gg,11005
 tests/unit/mcp/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -112,9 +116,9 @@ tests/unit/server/test_utils.py,sha256=kvhzdgNfsJl5tqcRBWg2yTR5GPpyrFCOmEIOuHb39
 tests/unit/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/utils/test_files.py,sha256=AKFmXQqXstyKd2PreE4EmQyhQYeqOmu1Sp80MwHrf_Q,5782
 tests/unit/vault/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/unit/vault/test_vault.py,sha256=R_RTDsralvE0JhuIrHQTb85gU9ipDJD7aoYIMdcg0o4,7264
-aixtools-0.1.8.dist-info/METADATA,sha256=aVW6zyI3cqprsWoTDOJCwXc-2tD2pMHb2bR6fc4Mo2w,18569
-aixtools-0.1.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.1.8.dist-info/entry_points.txt,sha256=dHoutULEZx7xXSqJrZdViSVjfInJibfLibi2nRXL3SE,56
-aixtools-0.1.8.dist-info/top_level.txt,sha256=ee4eF-0pqu45zCUVml0mWIhnXQgqMQper2-49BBVHLY,40
-aixtools-0.1.8.dist-info/RECORD,,
+tests/unit/vault/test_vault.py,sha256=T9V2Opxl3N5sJPftw0Q4lnVOs6urGpAmffe0cz6PUfw,10445
+aixtools-0.1.9.dist-info/METADATA,sha256=NyTMeH1c44cDNbXW7VKy2feT-xConc_uGIx0gSa68KM,18569
+aixtools-0.1.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.1.9.dist-info/entry_points.txt,sha256=dHoutULEZx7xXSqJrZdViSVjfInJibfLibi2nRXL3SE,56
+aixtools-0.1.9.dist-info/top_level.txt,sha256=ee4eF-0pqu45zCUVml0mWIhnXQgqMQper2-49BBVHLY,40
+aixtools-0.1.9.dist-info/RECORD,,

docker/mcp-base/Dockerfile

@@ -1,17 +1,14 @@
 FROM ubuntu:22.04
 
 RUN apt-get -y update && \
-    apt-get -y install git curl sudo ca-certificates
+    apt-get -y install ca-certificates curl gcc git libcap2-bin sudo
+RUN mv /usr/bin/sudo /usr/sbin
 
-# Build argument for custom certificate file (optional)
-ARG CUSTOM_CERT_FILE
-# Copy custom certificate if specified
-RUN if [ -n "$CUSTOM_CERT_FILE" ] && [ -f "$CUSTOM_CERT_FILE" ]; then \
-    cp "$CUSTOM_CERT_FILE" /usr/local/share/ca-certificates/ && \
-    update-ca-certificates && \
-    export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt && \
-    export CURL_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt; \
-    fi
+# Add Zscaler CA certificate
+COPY ./zscaler.crt /usr/local/share/ca-certificates/
+RUN update-ca-certificates
+ENV SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt
+ENV CURL_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt
 
 # Install `uv` Python package manager
 RUN bash -o pipefail -c "curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh"

docker/mcp-base/zscaler.crt

@@ -0,0 +1,28 @@
+-----BEGIN CERTIFICATE-----
+MIIE0zCCA7ugAwIBAgIJANu+mC2Jt3uTMA0GCSqGSIb3DQEBCwUAMIGhMQswCQYD
+VQQGEwJVUzETMBEGA1UECBMKQ2FsaWZvcm5pYTERMA8GA1UEBxMIU2FuIEpvc2Ux
+FTATBgNVBAoTDFpzY2FsZXIgSW5jLjEVMBMGA1UECxMMWnNjYWxlciBJbmMuMRgw
+FgYDVQQDEw9ac2NhbGVyIFJvb3QgQ0ExIjAgBgkqhkiG9w0BCQEWE3N1cHBvcnRA
+enNjYWxlci5jb20wHhcNMTQxMjE5MDAyNzU1WhcNNDIwNTA2MDAyNzU1WjCBoTEL
+MAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExETAPBgNVBAcTCFNhbiBK
+b3NlMRUwEwYDVQQKEwxac2NhbGVyIEluYy4xFTATBgNVBAsTDFpzY2FsZXIgSW5j
+LjEYMBYGA1UEAxMPWnNjYWxlciBSb290IENBMSIwIAYJKoZIhvcNAQkBFhNzdXBw
+b3J0QHpzY2FsZXIuY29tMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA
+qT7STSxZRTgEFFf6doHajSc1vk5jmzmM6BWuOo044EsaTc9eVEV/HjH/1DWzZtcr
+fTj+ni205apMTlKBW3UYR+lyLHQ9FoZiDXYXK8poKSV5+Tm0Vls/5Kb8mkhVVqv7
+LgYEmvEY7HPY+i1nEGZCa46ZXCOohJ0mBEtB9JVlpDIO+nN0hUMAYYdZ1KZWCMNf
+5J/aTZiShsorN2A38iSOhdd+mcRM4iNL3gsLu99XhKnRqKoHeH83lVdfu1XBeoQz
+z5V6gA3kbRvhDwoIlTBeMa5l4yRdJAfdpkbFzqiwSgNdhbxTHnYYorDzKfr2rEFM
+dsMU0DHdeAZf711+1CunuQIDAQABo4IBCjCCAQYwHQYDVR0OBBYEFLm33UrNww4M
+hp1d3+wcBGnFTpjfMIHWBgNVHSMEgc4wgcuAFLm33UrNww4Mhp1d3+wcBGnFTpjf
+oYGnpIGkMIGhMQswCQYDVQQGEwJVUzETMBEGA1UECBMKQ2FsaWZvcm5pYTERMA8G
+A1UEBxMIU2FuIEpvc2UxFTATBgNVBAoTDFpzY2FsZXIgSW5jLjEVMBMGA1UECxMM
+WnNjYWxlciBJbmMuMRgwFgYDVQQDEw9ac2NhbGVyIFJvb3QgQ0ExIjAgBgkqhkiG
+9w0BCQEWE3N1cHBvcnRAenNjYWxlci5jb22CCQDbvpgtibd7kzAMBgNVHRMEBTAD
+AQH/MA0GCSqGSIb3DQEBCwUAA4IBAQAw0NdJh8w3NsJu4KHuVZUrmZgIohnTm0j+
+RTmYQ9IKA/pvxAcA6K1i/LO+Bt+tCX+C0yxqB8qzuo+4vAzoY5JEBhyhBhf1uK+P
+/WVWFZN/+hTgpSbZgzUEnWQG2gOVd24msex+0Sr7hyr9vn6OueH+jj+vCMiAm5+u
+kd7lLvJsBu3AO3jGWVLyPkS3i6Gf+rwAp1OsRrv3WnbkYcFf9xjuaf4z0hRCrLN2
+xFNjavxrHmsH8jPHVvgc1VD0Opja0l/BRVauTrUaoW6tE+wFG5rEcPGS80jjHK4S
+pB5iDj2mUZH1T8lzYtuZy0ZPirxmtsk3135+CKNa2OCAhhFjE0xd
+-----END CERTIFICATE-----

scripts/test.sh

@@ -15,9 +15,16 @@ touch .env
 # Add the project root to PYTHONPATH so imports work correctly
 export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
 
-# Run all tests using pytest with coverage
-echo "Running tests with coverage..."
-pytest tests/ -v --cov=aixtools --cov-report=term-missing
-
-# Generate HTML coverage report
-# pytest tests/ --cov=aixtools --cov-report=html
+# Check if a specific test is provided as parameter
+if [ $# -eq 1 ]; then
+    # Run specific test without coverage
+    echo "Running specific test: $1"
+    pytest "$1" -v
+else
+    # Run all tests using pytest with coverage
+    echo "Running tests with coverage..."
+    pytest tests/ -v --cov=aixtools --cov-report=term-missing
+    
+    # Generate HTML coverage report
+    # pytest tests/ --cov=aixtools --cov-report=html
+fi

tests/unit/compliance/test_private_data.py

@@ -0,0 +1,329 @@
+import json
+import unittest
+from pathlib import Path
+
+from aixtools.compliance.private_data import PrivateData, PRIVATE_DATA_FILE
+from aixtools.server.path import get_workspace_path
+
+
+class TestPrivateData(unittest.TestCase):
+    """Test cases for PrivateData class without mocking or patching."""
+
+    def setUp(self):
+        """Set up test environment with temporary directory."""        
+        # Create a test user and session context
+        self.workspace_path = Path(get_workspace_path())
+        self.workspace_path.mkdir(parents=True, exist_ok=True)
+        self.private_data_file = self.workspace_path / PRIVATE_DATA_FILE
+
+    def tearDown(self):
+        """Clean up test environment."""
+        self.private_data_file.unlink(missing_ok=True)
+
+    def test_init_no_existing_file(self):
+        """Test initialization when no private data file exists."""
+        private_data = PrivateData()
+        self.assertFalse(private_data.has_private_data)
+        self.assertEqual(private_data.get_private_datasets(), [])
+        self.assertEqual(private_data.get_idap_datasets(), [])
+
+    def test_init_with_existing_file(self):
+        """Test initialization when private data file exists."""
+        # Create a test file
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        test_data = {
+            "_has_private_data": True,
+            "_private_datasets": ["dataset1", "dataset2"],
+            "_idap_datasets": ["dataset1"],
+            "ctx": None
+        }
+        with open(private_data_path, "w") as f:
+            json.dump(test_data, f)
+        
+        private_data = PrivateData()
+        
+        self.assertTrue(private_data.has_private_data)
+        self.assertEqual(set(private_data.get_private_datasets()), {"dataset1", "dataset2"})
+        self.assertEqual(private_data.get_idap_datasets(), ["dataset1"])
+
+    def test_add_private_dataset_new(self):
+        """Test adding a new private dataset."""
+        private_data = PrivateData()
+        
+        private_data.add_private_dataset("test_dataset")
+        
+        self.assertTrue(private_data.has_private_data)
+        self.assertIn("test_dataset", private_data.get_private_datasets())
+        
+        # Verify data is saved to file
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        self.assertTrue(private_data_path.exists())
+        
+        with open(private_data_path, "r") as f:
+            saved_data = json.load(f)
+        
+        self.assertTrue(saved_data["_has_private_data"])
+        self.assertIn("test_dataset", saved_data["_private_datasets"])
+
+    def test_add_private_dataset_duplicate(self):
+        """Test adding a duplicate private dataset."""
+        private_data = PrivateData()
+        
+        private_data.add_private_dataset("test_dataset")
+        initial_count = len(private_data.get_private_datasets())
+        
+        private_data.add_private_dataset("test_dataset")
+        final_count = len(private_data.get_private_datasets())
+        
+        self.assertEqual(initial_count, final_count)
+        self.assertEqual(private_data.get_private_datasets().count("test_dataset"), 1)
+
+    def test_add_idap_dataset_new(self):
+        """Test adding a new IDAP dataset."""
+        private_data = PrivateData()
+        
+        private_data.add_idap_dataset("idap_dataset")
+        
+        self.assertTrue(private_data.has_private_data)
+        self.assertIn("idap_dataset", private_data.get_idap_datasets())
+        self.assertIn("idap_dataset", private_data.get_private_datasets())
+        
+        # Verify data is saved to file
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        self.assertTrue(private_data_path.exists())
+        
+        with open(private_data_path, "r") as f:
+            saved_data = json.load(f)
+        
+        self.assertTrue(saved_data["_has_private_data"])
+        self.assertIn("idap_dataset", saved_data["_idap_datasets"])
+        self.assertIn("idap_dataset", saved_data["_private_datasets"])
+
+    def test_add_idap_dataset_existing_private(self):
+        """Test adding IDAP dataset when it already exists as private dataset."""
+        private_data = PrivateData()
+        
+        private_data.add_private_dataset("existing_dataset")
+        private_data.add_idap_dataset("existing_dataset")
+        
+        self.assertTrue(private_data.has_private_data)
+        self.assertIn("existing_dataset", private_data.get_idap_datasets())
+        self.assertIn("existing_dataset", private_data.get_private_datasets())
+        self.assertEqual(private_data.get_private_datasets().count("existing_dataset"), 1)
+
+    def test_add_idap_dataset_duplicate(self):
+        """Test adding a duplicate IDAP dataset."""
+        private_data = PrivateData()
+        
+        private_data.add_idap_dataset("idap_dataset")
+        initial_idap_count = len(private_data.get_idap_datasets())
+        initial_private_count = len(private_data.get_private_datasets())
+        
+        private_data.add_idap_dataset("idap_dataset")
+        final_idap_count = len(private_data.get_idap_datasets())
+        final_private_count = len(private_data.get_private_datasets())
+        
+        self.assertEqual(initial_idap_count, final_idap_count)
+        self.assertEqual(initial_private_count, final_private_count)
+
+    def test_has_private_dataset(self):
+        """Test checking if a private dataset exists."""
+        private_data = PrivateData()
+        
+        self.assertFalse(private_data.has_private_dataset("nonexistent"))
+        
+        private_data.add_private_dataset("test_dataset")
+        
+        self.assertTrue(private_data.has_private_dataset("test_dataset"))
+        self.assertFalse(private_data.has_private_dataset("nonexistent"))
+
+    def test_has_idap_dataset(self):
+        """Test checking if an IDAP dataset exists."""
+        private_data = PrivateData()
+        
+        self.assertFalse(private_data.has_idap_dataset("nonexistent"))
+        
+        private_data.add_idap_dataset("idap_dataset")
+        
+        self.assertTrue(private_data.has_idap_dataset("idap_dataset"))
+        self.assertFalse(private_data.has_idap_dataset("nonexistent"))
+
+    def test_get_datasets_returns_copies(self):
+        """Test that get methods return copies to prevent external modification."""
+        private_data = PrivateData()
+        private_data.add_private_dataset("private_dataset")
+        private_data.add_idap_dataset("idap_dataset")
+        
+        private_datasets = private_data.get_private_datasets()
+        idap_datasets = private_data.get_idap_datasets()
+        
+        # Modify the returned lists
+        private_datasets.append("external_modification")
+        idap_datasets.append("external_modification")
+        
+        # Verify original data is unchanged
+        self.assertNotIn("external_modification", private_data.get_private_datasets())
+        self.assertNotIn("external_modification", private_data.get_idap_datasets())
+
+    def test_has_private_data_setter_true(self):
+        """Test setting has_private_data to True."""
+        private_data = PrivateData()
+        
+        private_data.has_private_data = True
+        
+        self.assertTrue(private_data.has_private_data)
+        
+        # Verify data is saved to file
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        self.assertTrue(private_data_path.exists())
+
+    def test_has_private_data_setter_false(self):
+        """Test setting has_private_data to False clears all data."""
+        private_data = PrivateData()
+        private_data.add_private_dataset("dataset1")
+        private_data.add_idap_dataset("dataset2")
+        
+        private_data.has_private_data = False
+        
+        self.assertFalse(private_data.has_private_data)
+        self.assertEqual(private_data.get_private_datasets(), [])
+        self.assertEqual(private_data.get_idap_datasets(), [])
+        
+        # Verify file is deleted
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        self.assertFalse(private_data_path.exists())
+
+    def test_save_with_no_private_data(self):
+        """Test saving when has_private_data is False deletes the file."""
+        private_data = PrivateData()
+        private_data.add_private_dataset("test_dataset")
+        
+        # Verify file exists
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        self.assertTrue(private_data_path.exists())
+        
+        private_data.has_private_data = False
+        
+        # Verify file is deleted
+        self.assertFalse(private_data_path.exists())
+
+    def test_save_with_private_data(self):
+        """Test saving when has_private_data is True creates/updates the file."""
+        private_data = PrivateData()
+        private_data.add_private_dataset("dataset1")
+        private_data.add_idap_dataset("dataset2")
+        
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        self.assertTrue(private_data_path.exists())
+        
+        with open(private_data_path, "r") as f:
+            saved_data = json.load(f)
+        
+        self.assertTrue(saved_data["_has_private_data"])
+        self.assertIn("dataset1", saved_data["_private_datasets"])
+        self.assertIn("dataset2", saved_data["_private_datasets"])
+        self.assertIn("dataset2", saved_data["_idap_datasets"])
+        self.assertIsNone(saved_data["ctx"])
+
+    def test_load_from_corrupted_file(self):
+        """Test loading from a corrupted JSON file."""
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        with open(private_data_path, "w") as f:
+            f.write("invalid json content")
+        
+        with self.assertRaises(json.JSONDecodeError):
+            PrivateData()
+
+    def test_load_from_file_with_missing_fields(self):
+        """Test loading from a file with missing fields uses defaults."""
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        test_data = {"_has_private_data": True}  # Missing other fields
+        with open(private_data_path, "w") as f:
+            json.dump(test_data, f)
+        
+        private_data = PrivateData()
+        
+        self.assertTrue(private_data.has_private_data)
+        self.assertEqual(private_data.get_private_datasets(), [])
+        self.assertEqual(private_data.get_idap_datasets(), [])
+
+    def test_multiple_operations_persistence(self):
+        """Test that multiple operations are properly persisted."""
+        # Create first instance and add data
+        private_data1 = PrivateData()
+        private_data1.add_private_dataset("dataset1")
+        private_data1.add_idap_dataset("dataset2")
+        
+        # Create second instance and verify data is loaded
+        private_data2 = PrivateData()
+        self.assertTrue(private_data2.has_private_data)
+        self.assertIn("dataset1", private_data2.get_private_datasets())
+        self.assertIn("dataset2", private_data2.get_private_datasets())
+        self.assertIn("dataset2", private_data2.get_idap_datasets())
+        
+        # Add more data with second instance
+        private_data2.add_private_dataset("dataset3")
+        
+        # Create third instance and verify all data is present
+        private_data3 = PrivateData()
+        expected_private = {"dataset1", "dataset2", "dataset3"}
+        expected_idap = {"dataset2"}
+        
+        self.assertEqual(set(private_data3.get_private_datasets()), expected_private)
+        self.assertEqual(set(private_data3.get_idap_datasets()), expected_idap)
+
+    def test_str_and_repr(self):
+        """Test string representation methods."""
+        private_data = PrivateData()
+        private_data.add_private_dataset("dataset1")
+        private_data.add_idap_dataset("dataset2")
+        
+        str_repr = str(private_data)
+        repr_repr = repr(private_data)
+        
+        self.assertEqual(str_repr, repr_repr)
+        self.assertIn("has_private_data=True", str_repr)
+        self.assertIn("dataset1", str_repr)
+        self.assertIn("dataset2", str_repr)
+        self.assertIn(str(self.workspace_path / PRIVATE_DATA_FILE), str_repr)
+
+    def test_file_operations_create_directories(self):
+        """Test that file operations create necessary directories."""
+        # Remove the workspace directory
+        import shutil
+        shutil.rmtree(self.workspace_path)
+        self.assertFalse(self.workspace_path.exists())
+        
+        # Create PrivateData instance and add data
+        private_data = PrivateData()
+        private_data.add_private_dataset("test_dataset")
+        
+        # Verify directory and file were created
+        self.assertTrue(self.workspace_path.exists())
+        private_data_path = self.workspace_path / PRIVATE_DATA_FILE
+        self.assertTrue(private_data_path.exists())
+
+    def test_concurrent_modifications_simulation(self):
+        """Test simulation of concurrent modifications (without actual threading)."""
+        # Create two instances with same context
+        private_data1 = PrivateData()
+        private_data2 = PrivateData()
+        
+        # Add data with first instance
+        private_data1.add_private_dataset("dataset1")
+        
+        # Second instance should see the changes when reloaded
+        private_data2.load()
+        self.assertIn("dataset1", private_data2.get_private_datasets())
+        
+        # Add data with second instance
+        private_data2.add_idap_dataset("dataset2")
+        
+        # First instance should see the changes when reloaded
+        private_data1.load()
+        self.assertIn("dataset2", private_data1.get_idap_datasets())
+        self.assertIn("dataset2", private_data1.get_private_datasets())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/vault/test_vault.py

@@ -159,6 +159,81 @@ def test_read_user_service_secret_not_found(patched_vault_client):
     )
 
 
+def test_list_user_secret_keys_success(patched_vault_client):
+    """Test successful listing of all user secret keys."""
+    mock_response = {"data": {"keys": ["service1/", "service2/", "service3"]}}
+    patched_vault_client.client.secrets.kv.v2.list_secrets.return_value = mock_response
+
+    result = patched_vault_client.list_user_secret_keys(user_id="test-user")
+
+    assert result == ["service1", "service2", "service3"]
+    patched_vault_client.client.secrets.kv.v2.list_secrets.assert_called_once_with(
+        path="path/dev/test-user", mount_point="secret"
+    )
+
+
+def test_list_user_secret_keys_no_secrets(patched_vault_client):
+    """Test listing user secret keys when no secrets exist."""
+    mock_response = {"data": {"keys": []}}
+    patched_vault_client.client.secrets.kv.v2.list_secrets.return_value = mock_response
+
+    result = patched_vault_client.list_user_secret_keys(user_id="test-user")
+
+    assert result == []
+    patched_vault_client.client.secrets.kv.v2.list_secrets.assert_called_once_with(
+        path="path/dev/test-user", mount_point="secret"
+    )
+
+
+def test_list_user_secret_keys_empty_response(patched_vault_client):
+    """Test listing user secret keys when response is empty."""
+    patched_vault_client.client.secrets.kv.v2.list_secrets.return_value = {}
+
+    result = patched_vault_client.list_user_secret_keys(user_id="test-user")
+
+    assert result == []
+    patched_vault_client.client.secrets.kv.v2.list_secrets.assert_called_once_with(
+        path="path/dev/test-user", mount_point="secret"
+    )
+
+
+def test_list_user_secret_keys_invalid_path(patched_vault_client):
+    """Test listing user secret keys when user path does not exist."""
+    patched_vault_client.client.secrets.kv.v2.list_secrets.side_effect = InvalidPath
+
+    result = patched_vault_client.list_user_secret_keys(user_id="test-user")
+
+    assert result == []
+    patched_vault_client.client.secrets.kv.v2.list_secrets.assert_called_once_with(
+        path="path/dev/test-user", mount_point="secret"
+    )
+
+
+def test_list_user_secret_keys_unexpected_error(patched_vault_client):
+    """Test listing user secret keys when an unexpected exception occurs."""
+    patched_vault_client.client.secrets.kv.v2.list_secrets.side_effect = Exception("List error")
+
+    with pytest.raises(VaultAuthError, match="List error"):
+        patched_vault_client.list_user_secret_keys(user_id="test-user")
+
+    patched_vault_client.client.secrets.kv.v2.list_secrets.assert_called_once_with(
+        path="path/dev/test-user", mount_point="secret"
+    )
+
+
+def test_list_user_secret_keys_strips_trailing_slashes(patched_vault_client):
+    """Test that trailing slashes are properly stripped from directory names."""
+    mock_response = {"data": {"keys": ["service1/", "service2/", "service3", "service4/"]}}
+    patched_vault_client.client.secrets.kv.v2.list_secrets.return_value = mock_response
+
+    result = patched_vault_client.list_user_secret_keys(user_id="test-user")
+
+    assert result == ["service1", "service2", "service3", "service4"]
+    patched_vault_client.client.secrets.kv.v2.list_secrets.assert_called_once_with(
+        path="path/dev/test-user", mount_point="secret"
+    )
+
+
 def test_read_user_service_secret_error(patched_vault_client):
     """Test read_user_service_secret when an unexpected exception occurs."""
     patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = Exception("Read error")