Content
# QU3 - Quantum-Safe MCP Client
This project provides a client application (`qu3-app`) for secure interaction with Quantum-Safe Multi-Compute Provider (MCP) environments. It leverages post-quantum cryptography (PQC) standards for establishing secure communication channels, ensuring client authenticity, and verifying server attestations.
This client is designed to work with MCP servers that support the QU3 interaction protocols. For development and testing, a compatible mock server implementation is included in `scripts/mock_mcp_server.py`.
## Architecture & Flow
### Secure Communication Flow
The following diagram illustrates the end-to-end secure communication pattern implemented between the QU3 Client and the MCP Server:
```mermaid
sequenceDiagram
participant Client as QU3 Client (CLI)
participant Server as MCP Server
Note over Client,Server: Initial Setup (First Run / Keys Missing)
Client->>+Server: GET /keys (Fetch Server Public Keys)
Server-->>-Client: Server KEM PK, Server Sign PK (Base64)
Client->>Client: Store Server Public Keys Locally
Note over Client,Server: Establish Secure Session
Client->>+Server: POST /kem-handshake/initiate { Client KEM PK, Client Sign PK }
Server->>Server: Encapsulate Shared Secret (using Client KEM PK)
Server->>Server: Derive AES-256 Session Key (HKDF)
Server->>Server: Store Session { AES Key, Client Sign PK, Timestamp }
Server-->>-Client: { KEM Ciphertext }
Client->>Client: Decapsulate Shared Secret (using Client KEM SK)
Client->>Client: Derive AES-256 Session Key (HKDF)
Client->>Client: Store AES Session Key
Note over Client,Server: Secured Inference Request
Client->>Client: Prepare Payload { target_url, model_id, input_data }
Client->>Client: Sign(Payload) using Client Sign SK -> signature
Client->>Client: Prepare Encrypt Payload { Payload + signature }
Client->>Client: Encrypt(Encrypt Payload) using AES Key -> nonce, ciphertext
Client->>+Server: POST /inference { Client KEM PK, nonce, ciphertext }
Server->>Server: Lookup Session by Client KEM PK
Server->>Server: Validate Session Timestamp
Server->>Server: Decrypt(ciphertext) using AES Key -> payload+signature
Server->>Server: Verify(Payload, signature) using Client Sign PK
Server->>Server: Execute Model(input_data) -> output_data
Server->>Server: Prepare Attestation { metadata, status, hashes }
Server->>Server: Sign(Attestation) using Server Sign SK -> attestation_sig
Server->>Server: Prepare Response { status, output_data, attestation, attestation_sig }
Server->>Server: Encrypt(Response) using AES Key -> resp_nonce, resp_ciphertext
Server-->>-Client: { resp_nonce, resp_ciphertext }
Client->>Client: Decrypt(resp_ciphertext) using AES Key -> response
Client->>Client: Verify(Attestation, attestation_sig) using Server Sign PK
Client->>Client: Process Result
Note over Client,Server: Secured Policy Update
Client->>Client: Read Policy Content
Client->>Client: Sign(Policy) using Client Sign SK -> policy_sig
Client->>Client: Encrypt(Policy) using AES Key -> policy_nonce, policy_ciphertext
Client->>+Server: POST /policy-update { Client KEM PK, policy_nonce, policy_ciphertext, policy_sig }
Server->>Server: Lookup Session by Client KEM PK
Server->>Server: Validate Session Timestamp
Server->>Server: Decrypt(policy_ciphertext) using AES Key -> policy
Server->>Server: Verify(policy, policy_sig) using Client Sign PK
Server->>Server: Process Policy (Mock)
Server->>Server: Prepare Status Response { status }
Server->>Server: Sign(Status) using Server Sign SK -> status_sig
Server->>Server: Encrypt(Status) using AES Key -> status_nonce, status_ciphertext
Server-->>-Client: { status_nonce, status_ciphertext, status_sig }
Client->>Client: Decrypt(status_ciphertext) using AES Key -> status_response
Client->>Client: Verify(Status, status_sig) using Server Sign PK
Client->>Client: Display Status
```
### Client Component Interaction
This diagram shows how the main Python modules within the client application interact:
```mermaid
graph TD
A("User CLI (Typer)") --> B("src_main");
B -- Initiates --> C("src_mcp_client (MCPClient)");
B -- Uses --> D("src_config_utils");
C -- Uses --> E("src_pqc_utils");
C -- Uses --> F("src_models");
C -- Uses --> G("requests Session");
D -- Uses --> H("PyYAML");
D -- Manages --> I("Key Files");
D -- Uses --> G;
E -- Uses --> J("liboqs_python");
E -- Uses --> K("cryptography");
subgraph Cryptography
J
K
end
subgraph Networking
G
end
subgraph "Configuration & Keys"
H
I
end
```
### Key Management Overview
Keys are crucial for the security protocols. Here's how they are managed:
```mermaid
graph LR
subgraph Client Side
A["CLI: generate-keys"] --> B{"src_main.py"};
B --> C["src_pqc_utils.py"]:::pqc --> D{"Generate KEM/Sign Pairs"};
D --> E["src_config_utils.py"]:::config --> F["Save Client Keys (.pub, .sec)"];
G["CLI: run-inference/etc."] --> B;
B --> H{"Initialize Client"};
H --> E --> I{"Load Client Keys?"};
I -- Found --> J["Use Keys"];
I -- Not Found --> D;
H --> E --> K{"Load Server Keys?"};
K -- Found --> J;
K -- Not Found --> L{"Fetch Server Keys?"};
L -- Calls --> E --> M["GET /keys"]:::net;
M -- Response --> E --> N["Save Server Keys (.pub)"];
N --> J;
L -- Fetch Fail --> O["(Error - Cannot Proceed)"];
end
subgraph "Server Side (Mock)"
P["Server Startup"] --> Q["scripts_mock_mcp_server.py"];
Q --> R["src_config_utils.py"]:::config --> S{"Load/Generate Server Keys"};
S --> T["Save Server Keys (.pub, .sec)"];
Q --> U["Register /keys Endpoint"];
U -- Request for /keys --> V{"Return Server Public Keys"};
end
subgraph "Filesystem (Key Dir)"
F
T
N
end
classDef pqc fill:#f9d,stroke:#333,stroke-width:2px;
classDef config fill:#cfc,stroke:#333,stroke-width:2px;
classDef net fill:#cdf,stroke:#333,stroke-width:2px;
```
## Core Components
* **`src/main.py`**: Command-line interface (CLI) built with Typer. Handles user commands, orchestrates client operations, and displays results. Includes commands for key generation, inference, agent workflows, and policy updates.
* **`src/mcp_client.py`**: The main client class (`MCPClient`) responsible for:
* Managing PQC keys.
* Establishing secure sessions via KEM handshake (`connect`).
* Sending signed and encrypted requests (`send_request`).
* Processing and verifying encrypted/signed responses.
* Handling disconnection (`disconnect`).
* **`src/pqc_utils.py`**: Utility functions for PQC operations (Kyber KEM, SPHINCS+ signing) using `liboqs-python`, AES-GCM encryption/decryption using `cryptography`, and HKDF key derivation.
* **`src/config_utils.py`**: Handles loading configuration from `config.yaml`, loading/saving keys from/to files, and fetching server public keys from the `/keys` endpoint.
* **`src/models.py`**: Defines Pydantic models (`MCPRequest`, `MCPResponse`) for data structures exchanged between client components.
* **`scripts/mock_mcp_server.py`**: A FastAPI development/test server that simulates an MCP environment. Implements the server-side logic for KEM handshake, request decryption/verification, basic model execution, attestation signing, response encryption, policy updates, and key distribution.
* **`config.yaml`**: Configuration file for storing settings like the default key directory (`key_directory`) and server URL (`server_url`).
* **`tests/`**: Directory containing unit tests (`unittest`) for core components (`pqc_utils`, `config_utils`, `mcp_client`).
## Features
* **PQC Algorithms**: Uses NIST PQC finalists:
* KEM: `Kyber-768`
* Signature: `SPHINCS+-SHA2-128f-simple`
* **PQC Key Management**: Generates and loads Kyber and SPHINCS+ key pairs.
* **Secure Session Establishment**: Uses Kyber KEM over a network handshake (`/kem-handshake/initiate`) to establish a shared secret.
* **Key Derivation**: Derives a 32-byte AES-256 key from the KEM shared secret using HKDF-SHA256.
* **Encrypted Communication**: Encrypts request/response payloads (after KEM handshake) using AES-256-GCM with the derived session key.
* **Client Authentication**: Client signs requests using SPHINCS+; server verifies.
* **Server Attestation**: Server signs responses (attestation data) using SPHINCS+; client verifies.
* **Configuration**: Loads key directory and server URL from `config.yaml`.
* **Automated Server Key Fetching**: Client automatically fetches server public keys from the `/keys` endpoint if not found locally.
* **CLI Commands**:
* `generate-keys`: Creates client key pairs.
* `run-inference`: Sends a single, secured inference request.
* `run-agent`: Executes a sequential workflow (`modelA->modelB`), passing outputs as inputs (wraps non-dict output), with step-by-step reporting and robust failure handling.
* `update-policy`: Sends an encrypted and signed policy file to the server.
* **Mock Server**: Includes endpoints (`/`, `/keys`, `/kem-handshake/initiate`, `/inference`, `/policy-update`) implementing the corresponding server-side PQC and communication logic for testing. Provides example models `model_caps` and `model_reverse`.
* **Unit Tests**: Includes unit tests covering core cryptographic utilities, configuration management, and client communication logic (with network mocking).
## Setup
1. **Clone the repository:**
```bash
git clone <repository_url>
cd qu3-app
```
2. **Create a virtual environment:**
```bash
python3 -m venv venv
source venv/bin/activate
```
3. **Install dependencies:**
```bash
pip install -r requirements.txt
```
*(Note: `liboqs-python` might require system dependencies like a C compiler and the `liboqs` C library. Refer to its documentation if installation fails.)*
4. **(Optional) Configure `config.yaml`:** Modify `key_directory` or `server_url` if needed. The default key directory is `~/.qu3/keys/`.
## Running the Mock Server
In one terminal, run:
```bash
# Ensure virtual environment is active
source venv/bin/activate
python -m scripts.mock_mcp_server
```
The server will start (usually on `http://127.0.0.1:8000`) and automatically generate its own key pairs in the configured key directory if they don't exist.
## Running the Client CLI
In another terminal (with the virtual environment activated):
1. **Generate Client Keys:** (Only needs to be done once unless `--force` is used)
```bash
python -m src.main generate-keys
```
2. **Fetch Server Keys (Automatic):** The client will attempt to fetch keys from the server's `/keys` endpoint during initialization (`run-inference`, `run-agent`, `update-policy`) if `server_kem.pub` and `server_sign.pub` are not found in the key directory specified in `config.yaml`. Ensure the mock server is running before executing client commands that require connection.
3. **Run Single Inference:**
```bash
# Example using mock server's model_caps
python -m src.main run-inference model_caps '{"text": "process this data"}'
# Example specifying server URL
python -m src.main run-inference model_reverse '{"text": "backward"}' --server-url http://127.0.0.1:8000
```
4. **Run Agent Workflow:**
```bash
# Example chaining two mock models
python -m src.main run-agent "model_caps -> model_reverse" '{"text": "flow start"}'
```
5. **Update Policy:**
Create a policy file (e.g., `my_policy.txt`) with some content.
```bash
echo "Allow model_caps access." > my_policy.txt
python -m src.main update-policy --policy-file my_policy.txt
```
## Running Tests
To run the unit tests:
```bash
# Ensure virtual environment is active
source venv/bin/activate
python -m unittest discover tests
```
Or run specific test files:
```bash
python -m unittest tests.test_pqc_utils
python -m unittest tests.test_config_utils
python -m unittest tests.test_mcp_client
```
