feat(attestation): add vtpm attestation module

pyproject.toml

@@ -5,12 +5,15 @@ description = "Flare AI Kit template for single-node multi-model Consensus Learn
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
+    "cryptography>=44.0.1",
     "httpx>=0.28.1",
     "matplotlib>=3.10.0",
     "numpy>=2.2.2",
     "openai>=1.61.0",
     "pandas>=2.2.3",
     "pydantic-settings>=2.7.1",
+    "pyjwt>=2.10.1",
+    "pyopenssl>=25.0.0",
     "python-dotenv>=1.0.1",
     "requests>=2.32.3",
     "structlog>=25.1.0",

src/flare_ai_consensus/attestation/__init__.py

@@ -0,0 +1,21 @@
+from .vtpm_attestation import (
+    Vtpm,
+    VtpmAttestationError,
+)
+from .vtpm_validation import (
+    CertificateParsingError,
+    InvalidCertificateChainError,
+    SignatureValidationError,
+    VtpmValidation,
+    VtpmValidationError,
+)
+
+__all__ = [
+    "CertificateParsingError",
+    "InvalidCertificateChainError",
+    "SignatureValidationError",
+    "Vtpm",
+    "VtpmAttestationError",
+    "VtpmValidation",
+    "VtpmValidationError",
+]

src/flare_ai_consensus/attestation/simulated_token.txt

@@ -0,0 +1 @@
+eyJhbGciOiJSUzI1NiIsImtpZCI6IjQ2NzZjNDkwZGM0MzgyOTYzNjU5NTQ0MmU5M2NkYzVkMjdhYWEyNzkiLCJ0eXAiOiJKV1QifQ.eyJhdWQiOiJodHRwczovL3N0cy5nb29nbGUuY29tIiwiZXhwIjoxNzMwNjgxNjEyLCJpYXQiOjE3MzA2NzgwMTIsImlzcyI6Imh0dHBzOi8vY29uZmlkZW50aWFsY29tcHV0aW5nLmdvb2dsZWFwaXMuY29tIiwibmJmIjoxNzMwNjc4MDEyLCJzdWIiOiJodHRwczovL3d3dy5nb29nbGVhcGlzLmNvbS9jb21wdXRlL3YxL3Byb2plY3RzL2ZsYXJlLW5ldHdvcmstc2FuZGJveC96b25lcy91cy1jZW50cmFsMS1iL2luc3RhbmNlcy90ZXN0LWNvbmZpZGVudGlhbCIsImVhdF9ub25jZSI6IjB4MDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwZEVhRCIsImVhdF9wcm9maWxlIjoiaHR0cHM6Ly9jbG91ZC5nb29nbGUuY29tL2NvbmZpZGVudGlhbC1jb21wdXRpbmcvY29uZmlkZW50aWFsLXNwYWNlL2RvY3MvcmVmZXJlbmNlL3Rva2VuLWNsYWltcyIsInNlY2Jvb3QiOnRydWUsIm9lbWlkIjoxMTEyOSwiaHdtb2RlbCI6IkdDUF9BTURfU0VWIiwic3duYW1lIjoiQ09ORklERU5USUFMX1NQQUNFIiwic3d2ZXJzaW9uIjpbIjI0MDkwMCJdLCJkYmdzdGF0IjoiZW5hYmxlZCIsInN1Ym1vZHMiOnsiY29uZmlkZW50aWFsX3NwYWNlIjp7Im1vbml0b3JpbmdfZW5hYmxlZCI6eyJtZW1vcnkiOmZhbHNlfX0sImNvbnRhaW5lciI6eyJpbWFnZV9yZWZlcmVuY2UiOiJnaGNyLmlvL2RpbmVzaHBpbnRvL3Rlc3QtY29uZmlkZW50aWFsOm1haW4iLCJpbWFnZV9kaWdlc3QiOiJzaGEyNTY6YWY3MzhmZGRkMzFlYmU0OGVkNGQ4ZWM5MzZmMjQyMTI3ODUwZDZiMDFhYTY4YTFmYzliZGViOWMwNjNmZWI3YyIsInJlc3RhcnRfcG9saWN5IjoiTmV2ZXIiLCJpbWFnZV9pZCI6InNoYTI1Njo3YmNiMGU5OWUwOGMzNGM3NTkxOWE2NTQwYzJhMDdjMWRkODQ0Y2QzN2Y5MjdjZjBmNTg4ZjBkMzVmYzZmMWViIiwiZW52X292ZXJyaWRlIjp7IkFVRElFTkNFIjoiaHR0cHM6Ly9zdHMuZ29vZ2xlLmNvbSIsIk5PTkNFIjoiMHgwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDBkRWFEIn0sImVudiI6eyJBVURJRU5DRSI6Imh0dHBzOi8vc3RzLmdvb2dsZS5jb20iLCJHUEdfS0VZIjoiNzE2OTYwNUY2MkM3NTEzNTZEMDU0QTI2QTgyMUU2ODBFNUZBNjMwNSIsIkhPU1ROQU1FIjoidGVzdC1jb25maWRlbnRpYWwiLCJMQU5HIjoiQy5VVEYtOCIsIk5PTkNFIjoiMHgwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDBkRWFEIiwiUEFUSCI6Ii91c3IvbG9jYWwvYmluOi91c3IvbG9jYWwvc2JpbjovdXNyL2xvY2FsL2JpbjovdXNyL3NiaW46L3Vzci9iaW46L3NiaW46L2JpbiIsIlBZVEhPTl9TSEEyNTYiOiIyNDg4N2I5MmUyYWZkNGEyYWM2MDI0MTlhZDRiNTk2MzcyZjY3YWM5YjA3NzE5MGY0NTlhYmEzOTBmYWY1NTUwIiwiUFlUSE9OX1ZFUlNJT04iOiIzLjEyLjcifSwiYXJncyI6WyJ1diIsInJ1biIsImF0dGVzdGF0aW9uLnB5Il19LCJnY2UiOnsiem9uZSI6InVzLWNlbnRyYWwxLWIiLCJwcm9qZWN0X2lkIjoiZmxhcmUtbmV0d29yay1zYW5kYm94IiwicHJvamVjdF9udW1iZXIiOiI4MzY3NDUxNzg3NiIsImluc3RhbmNlX25hbWUiOiJ0ZXN0LWNvbmZpZGVudGlhbCIsImluc3RhbmNlX2lkIjoiMjMwOTAyNDU0MzcxMDQ5MzQ4NyJ9fSwiZ29vZ2xlX3NlcnZpY2VfYWNjb3VudHMiOlsiODM2NzQ1MTc4NzYtY29tcHV0ZUBkZXZlbG9wZXIuZ3NlcnZpY2VhY2NvdW50LmNvbSJdfQ.f2VAbbNl1N9CvL69HJzNKzqdxo4xVK9IVBaO7Wwp0gD8L8IKrvqSUzzXE_guw3hpX2enEnTUEzL6PqLj0bvCB8lMcwogKvhnV2q-WgOSHn3kPMZthrnTXtNarIOqZFTFty3HkFNjCRoE2isosS4rf9QLgASA5C4ASFGUUuFZhODC68sAWTB8mGkd4qTORF8yy5-2i_JgOCZVQhKKJLaEXwvUZmJXYO5i2OkkcFSoYnS1YvfobFi8zuiRIpqx-cv5aDGI6i91iXjk42Ljc4-7BYV_gLsf-p3lBvcEq9es-dGFUTUHLeUmhBXdpRaSgRgWkOgF6XNoLl4movIBZwLgvA

src/flare_ai_consensus/attestation/vtpm_attestation.py

@@ -0,0 +1,146 @@
+"""
+Client for communicating with the Confidential Space vTPM attestation service.
+
+This module provides a client to request attestation tokens from a local Unix domain
+socket endpoint. It extends HTTPConnection to handle Unix socket communication and
+implements token request functionality with nonce validation.
+
+Classes:
+    VtpmAttestationError: Exception for attestation service communication errors
+    VtpmAttestation: Client for requesting attestation tokens
+"""
+
+import json
+import socket
+from http.client import HTTPConnection
+from pathlib import Path
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+def get_simulated_token() -> str:
+    """Reads the first line from a given file path."""
+    with (Path(__file__).parent / "simulated_token.txt").open("r") as f:
+        return f.readline().strip()
+
+
+SIM_TOKEN = get_simulated_token()
+
+
+class VtpmAttestationError(Exception):
+    """
+    Exception raised for attestation service communication errors.
+
+    This includes invalid nonce values, connection failures, and
+    unexpected responses from the attestation service.
+    """
+
+
+class Vtpm:
+    """
+    Client for requesting attestation tokens via Unix domain socket."""
+
+    def __init__(
+        self,
+        url: str = "http://localhost/v1/token",
+        unix_socket_path: str = "/run/container_launcher/teeserver.sock",
+        simulate: bool = False,  # noqa: FBT001, FBT002
+    ) -> None:
+        self.url = url
+        self.unix_socket_path = unix_socket_path
+        self.simulate = simulate
+        self.attestation_requested: bool = False
+        self.logger = logger.bind(router="vtpm")
+        self.logger.debug(
+            "vtpm", simulate=simulate, url=url, unix_socket_path=self.unix_socket_path
+        )
+
+    def _check_nonce_length(self, nonces: list[str]) -> None:
+        """
+        Validate the byte length of provided nonces.
+
+        Each nonce must be between 10 and 74 bytes when UTF-8 encoded.
+
+        Args:
+            nonces: List of nonce strings to validate
+
+        Raises:
+            VtpmAttestationError: If any nonce is outside the valid length range
+        """
+        min_byte_len = 10
+        max_byte_len = 74
+        for nonce in nonces:
+            byte_len = len(nonce.encode("utf-8"))
+            self.logger.debug("nonce_length", byte_len=byte_len)
+            if byte_len < min_byte_len or byte_len > max_byte_len:
+                msg = f"Nonce '{nonce}' must be between {min_byte_len} bytes"
+                f" and {max_byte_len} bytes"
+                raise VtpmAttestationError(msg)
+
+    def get_token(
+        self,
+        nonces: list[str],
+        audience: str = "https://sts.google.com",
+        token_type: str = "OIDC",  # noqa: S107
+    ) -> str:
+        """
+        Request an attestation token from the service.
+
+        Requests a token with specified nonces for replay protection,
+        targeted at the specified audience. Supports both OIDC and PKI
+        token types.
+
+        Args:
+            nonces: List of random nonce strings for replay protection
+            audience: Intended audience for the token (default: "https://sts.google.com")
+            token_type: Type of token, either "OIDC" or "PKI" (default: "OIDC")
+
+        Returns:
+            str: The attestation token in JWT format
+
+        Raises:
+            VtpmAttestationError: If token request fails for any reason
+                (invalid nonces, service unavailable, etc.)
+
+        Example:
+            client = VtpmAttestation()
+            token = client.get_token(
+                nonces=["random_nonce"],
+                audience="https://my-service.example.com",
+                token_type="OIDC"
+            )
+        """
+        self._check_nonce_length(nonces)
+        if self.simulate:
+            self.logger.debug("sim_token", token=SIM_TOKEN)
+            return SIM_TOKEN
+
+        # Connect to the socket
+        client_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        client_socket.connect(self.unix_socket_path)
+
+        # Create an HTTP connection object
+        conn = HTTPConnection("localhost", timeout=10)
+        conn.sock = client_socket
+
+        # Send a POST request
+        headers = {"Content-Type": "application/json"}
+        body = json.dumps(
+            {"audience": audience, "token_type": token_type, "nonces": nonces}
+        )
+        conn.request("POST", self.url, body=body, headers=headers)
+
+        # Get and decode the response
+        res = conn.getresponse()
+        success_status = 200
+        if res.status != success_status:
+            msg = f"Failed to get attestation response: {res.status} {res.reason}"
+            raise VtpmAttestationError(msg)
+        token = res.read().decode()
+        self.logger.debug("token", token_type=token_type, token=token)
+
+        # Close the connection
+        conn.close()
+        return token