added code to store encrypted secrets

Author: NiveditJain

state-manager/app/controller/upsert_graph_template.py

@@ -18,7 +18,7 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
                 namespace_name=namespace_name,
                 x_exosphere_request_id=x_exosphere_request_id)
 
-            await graph_template.update(
+            await graph_template.set_secrets(body.secrets).update(
                 Set({
                     GraphTemplate.nodes: body.nodes, # type: ignore
                     GraphTemplate.validation_status: GraphTemplateValidationStatus.PENDING, # type: ignore
@@ -40,15 +40,14 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
                     nodes=body.nodes,
                     validation_status=GraphTemplateValidationStatus.PENDING,
                     validation_errors=[]
-                )
+                ).set_secrets(body.secrets)
             )
 
         return UpsertGraphTemplateResponse(
-            name=graph_template.name,
-            namespace=graph_template.namespace,
             nodes=graph_template.nodes,
             validation_status=graph_template.validation_status,
             validation_errors=graph_template.validation_errors,
+            secrets={secret_name: True for secret_name in graph_template.get_secrets().keys()},
             created_at=graph_template.created_at,
             updated_at=graph_template.updated_at
         )

state-manager/app/models/db/graph_template_model.py

@@ -1,17 +1,22 @@
+import re
+import base64
+
 from .base import BaseDatabaseModel
-from pydantic import Field
+from pydantic import Field, field_validator
 from typing import Optional, List
 from ..graph_template_validation_status import GraphTemplateValidationStatus
 from ..node_template_model import NodeTemplate
 from pymongo import IndexModel
-
+from typing import Dict
+from app.utils.encrypter import encrypter
 
 class GraphTemplate(BaseDatabaseModel):
     name: str = Field(..., description="Name of the graph")
     namespace: str = Field(..., description="Namespace of the graph")
     nodes: List[NodeTemplate] = Field(..., description="Nodes of the graph")
     validation_status: GraphTemplateValidationStatus = Field(..., description="Validation status of the graph")
     validation_errors: Optional[List[str]] = Field(None, description="Validation errors of the graph")
+    secrets: Dict[str, str] = Field(default_factory=dict, description="Secrets of the graph")
 
     class Settings:
         indexes = [
@@ -20,4 +25,58 @@ class Settings:
                 unique=True,
                 name="unique_name_namespace"
             )
-        ]
+        ]
+
+    @field_validator('secrets')
+    @classmethod
+    def validate_secrets(cls, v: Dict[str, str]) -> Dict[str, str]:
+        for secret_name, secret_value in v.items():
+            if not secret_name or not secret_value:
+                raise ValueError("Secrets cannot be empty")
+            if not isinstance(secret_name, str):
+                raise ValueError("Secret name must be a string")
+            if not isinstance(secret_value, str):
+                raise ValueError("Secret value must be a string")
+            cls._validate_secret_value(secret_value)
+            
+        return v
+    
+    @classmethod
+    def _validate_secret_value(cls, secret_value: str) -> None:
+        # Check minimum length for AES-GCM encrypted string
+        # 12 bytes nonce + minimum ciphertext + base64 encoding
+        if len(secret_value) < 32:  # Minimum length for encrypted string
+            raise ValueError("Value appears to be too short for an encrypted string")
+        
+        # Check if the string contains only URL-safe base64 characters
+        url_safe_base64_pattern = r'^[A-Za-z0-9_-]+$'
+        if not re.match(url_safe_base64_pattern, secret_value):
+            raise ValueError("Value must be URL-safe base64 encoded")
+        
+        # Check if the string length is valid for base64 encoding
+        # Base64 encoding increases size by ~33%, and we need at least 12 bytes nonce
+        if len(secret_value) % 4 != 0:
+            raise ValueError("Value length is not valid for base64 encoding")
+        
+        # Try to decode as base64 to ensure it's valid
+        try:
+            decoded = base64.urlsafe_b64decode(secret_value)
+            if len(decoded) < 12:
+                raise ValueError("Decoded value is too short to contain valid nonce")
+        except Exception:
+            raise ValueError("Value is not valid URL-safe base64 encoded")
+        
+
+    def set_secrets(self, secrets: Dict[str, str]) -> "GraphTemplate":
+        self.secrets = {secret_name: encrypter.encrypt(secret_value) for secret_name, secret_value in secrets.items()}
+        return self
+    
+    def get_secrets(self) -> Dict[str, str]:
+        if not self.secrets:
+            return {}
+        return {secret_name: encrypter.decrypt(secret_value) for secret_name, secret_value in self.secrets.items()}
+    
+    def get_secret(self, secret_name: str) -> str:
+        if not self.secrets:
+            return ""
+        return encrypter.decrypt(self.secrets[secret_name])

state-manager/app/models/graph_models.py

@@ -1,16 +1,18 @@
 from .node_template_model import NodeTemplate
 from pydantic import BaseModel, Field
-from typing import List, Optional
+from typing import Dict, List, Optional
 from datetime import datetime
 from .graph_template_validation_status import GraphTemplateValidationStatus
 
 
 class UpsertGraphTemplateRequest(BaseModel):
+    secrets: Dict[str, str] = Field(..., description="Dictionary of secrets that are used while graph execution")
     nodes: List[NodeTemplate] = Field(..., description="List of node templates that define the graph structure")
 
 
 class UpsertGraphTemplateResponse(BaseModel):
     nodes: List[NodeTemplate] = Field(..., description="List of node templates that define the graph structure")
+    secrets: Dict[str, bool] = Field(..., description="Dictionary of secrets that are used while graph execution")
     created_at: datetime = Field(..., description="Timestamp when the graph template was created")
     updated_at: datetime = Field(..., description="Timestamp when the graph template was last updated")
     validation_status: GraphTemplateValidationStatus = Field(..., description="Current validation status of the graph template")

state-manager/app/utils/encrypter.py

@@ -0,0 +1,34 @@
+import os
+import base64
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+class Encrypter:
+
+    @staticmethod
+    def generate_key() -> str:
+        return base64.urlsafe_b64encode(AESGCM.generate_key(bit_length=256)).decode()
+
+    def __init__(self):
+        key_b64 = os.getenv("SECRETS_ENCRYPTION_KEY")
+        if not key_b64:
+            raise ValueError("SECRETS_ENCRYPTION_KEY is not set")
+        try:
+            self._key = base64.urlsafe_b64decode(key_b64)
+        except Exception as exc:
+            raise ValueError("Key must be URL-safe base64 (44 chars for 32-byte key)") from exc
+        if len(self._key) != 32:
+            raise ValueError("Key must be 32 raw bytes (256 bits)")
+        self._aesgcm = AESGCM(self._key)
+    
+    def encrypt(self, secret: str) -> str:
+        nonce = os.urandom(12)
+        ciphertext = self._aesgcm.encrypt(nonce, secret.encode(), None)
+        return base64.urlsafe_b64encode(nonce + ciphertext).decode()
+    
+    def decrypt(self, encrypted_secret: str) -> str:
+        encrypted_secret_bytes = base64.urlsafe_b64decode(encrypted_secret)
+        nonce = encrypted_secret_bytes[:12]
+        ciphertext = encrypted_secret_bytes[12:]
+        return self._aesgcm.decrypt(nonce, ciphertext, None).decode()
+
+encrypter = Encrypter()

state-manager/pyproject.toml

@@ -6,6 +6,7 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "beanie>=2.0.0",
+    "cryptography>=45.0.5",
     "fastapi>=0.116.1",
     "python-dotenv>=1.1.1",
     "structlog>=25.4.0",