auth

backend/src/agent/tools.py

@@ -572,8 +572,8 @@ async def analyze_chart_data(python_script: str, countback: Optional[int] = None
     plot_urls = []
 
     # Determine uploads directory (relative to this file)
-    uploads_dir = Path(__file__).parent.parent.parent / "uploads"
-    uploads_dir.mkdir(exist_ok=True)
+    uploads_dir = Path(__file__).parent.parent.parent / "data" / "uploads"
+    uploads_dir.mkdir(parents=True, exist_ok=True)
 
     try:
         with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):

backend/src/main.py

@@ -14,7 +14,7 @@ from pydantic import BaseModel
 import uuid
 import shutil
 
-from sync.protocol import HelloMessage, PatchMessage
+from sync.protocol import HelloMessage, PatchMessage, AuthMessage, AuthResponseMessage
 from sync.registry import SyncRegistry
 from gateway.hub import Gateway
 from gateway.channels.websocket import WebSocketChannel
@@ -26,6 +26,7 @@ from schema.chart_state import ChartState
 from datasource.registry import DataSourceRegistry
 from datasource.subscription_manager import SubscriptionManager
 from datasource.websocket_handler import DatafeedWebSocketHandler
+from secrets_manager import SecretsStore, InvalidMasterPassword
 
 # Configure logging
 logging.basicConfig(
@@ -52,6 +53,9 @@ agent_executor = None
 datasource_registry = DataSourceRegistry()
 subscription_manager = SubscriptionManager()
 
+# Global secrets store
+secrets_store = SecretsStore()
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -76,8 +80,19 @@ async def lifespan(app: FastAPI):
         logger.warning(f"CCXT not available: {e}. Only demo source will be available.")
         logger.info("To use real exchange data, install ccxt: pip install ccxt>=4.0.0")
 
-    # Get API keys from environment
-    anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
+    # Get API keys from secrets store if unlocked, otherwise fall back to environment
+    anthropic_api_key = None
+
+    if secrets_store.is_unlocked:
+        anthropic_api_key = secrets_store.get("ANTHROPIC_API_KEY")
+        if anthropic_api_key:
+            logger.info("Loaded API key from encrypted secrets store")
+
+    # Fall back to environment variable
+    if not anthropic_api_key:
+        anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
+        if anthropic_api_key:
+            logger.info("Loaded API key from environment")
 
     if not anthropic_api_key:
         logger.error("ANTHROPIC_API_KEY not found in environment!")
@@ -117,8 +132,8 @@ async def lifespan(app: FastAPI):
 app = FastAPI(lifespan=lifespan)
 
 # Create uploads directory
-UPLOAD_DIR = Path(__file__).parent.parent / "uploads"
-UPLOAD_DIR.mkdir(exist_ok=True)
+UPLOAD_DIR = Path(__file__).parent.parent / "data" / "uploads"
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
 
 # Mount static files for serving uploads
 app.mount("/uploads", StaticFiles(directory=str(UPLOAD_DIR)), name="uploads")
@@ -142,14 +157,6 @@ registry.register(chart_store, store_name="ChartStore")
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
-    registry.websocket = websocket
-
-    # Create WebSocket channel for agent communication
-    channel_id = f"ws_{id(websocket)}"
-    client_id = f"client_{id(websocket)}"
-    logger.info(f"WebSocket connected - channel_id: {channel_id}, client_id: {client_id}")
-    ws_channel = WebSocketChannel(channel_id, websocket, session_id="default")
-    gateway.register_channel(ws_channel)
 
     # Helper function to send responses
     async def send_response(response):
@@ -158,6 +165,168 @@ async def websocket_endpoint(websocket: WebSocket):
         except Exception as e:
             logger.error(f"Error sending response: {e}")
 
+    # Authentication state
+    is_authenticated = False
+
+    # Wait for authentication message (must be first message)
+    try:
+        auth_timeout = 30  # 30 seconds to authenticate
+        auth_data = await asyncio.wait_for(websocket.receive_text(), timeout=auth_timeout)
+        auth_message_json = json.loads(auth_data)
+
+        if auth_message_json.get("type") != "auth":
+            logger.warning("First message was not auth message")
+            await send_response(AuthResponseMessage(
+                success=False,
+                message="First message must be authentication"
+            ))
+            await websocket.close(code=1008, reason="Authentication required")
+            return
+
+        auth_msg = AuthMessage(**auth_message_json)
+        logger.info("Received authentication message")
+
+        # Check if secrets store needs initialization
+        if not secrets_store.is_initialized:
+            logger.info("Secrets store not initialized, performing first-time setup")
+
+            # Require password confirmation for initialization
+            if not auth_msg.confirm_password:
+                await send_response(AuthResponseMessage(
+                    success=False,
+                    needs_confirmation=True,
+                    message="First-time setup: password confirmation required"
+                ))
+                await websocket.close(code=1008, reason="Password confirmation required")
+                return
+
+            if auth_msg.password != auth_msg.confirm_password:
+                await send_response(AuthResponseMessage(
+                    success=False,
+                    needs_confirmation=True,
+                    message="Passwords do not match"
+                ))
+                await websocket.close(code=1008, reason="Password confirmation failed")
+                return
+
+            # Initialize secrets store
+            try:
+                secrets_store.initialize(auth_msg.password)
+
+                # Migrate ANTHROPIC_API_KEY from environment if present
+                env_key = os.environ.get("ANTHROPIC_API_KEY")
+                if env_key:
+                    secrets_store.set("ANTHROPIC_API_KEY", env_key)
+                    logger.info("Migrated ANTHROPIC_API_KEY from environment to secrets store")
+
+                is_authenticated = True
+                await send_response(AuthResponseMessage(
+                    success=True,
+                    message="Secrets store initialized successfully"
+                ))
+                logger.info("Secrets store initialized and authenticated")
+            except Exception as e:
+                logger.error(f"Failed to initialize secrets store: {e}")
+                await send_response(AuthResponseMessage(
+                    success=False,
+                    message=f"Initialization failed: {str(e)}"
+                ))
+                await websocket.close(code=1011, reason="Initialization failed")
+                return
+        else:
+            # Unlock existing secrets store (or verify password if already unlocked)
+            try:
+                # If already unlocked, just verify the password is correct
+                if secrets_store.is_unlocked:
+                    # Verify password by creating a temporary store and attempting unlock
+                    from secrets_manager import SecretsStore as TempStore
+                    temp_store = TempStore(data_dir=secrets_store.data_dir)
+                    temp_store.unlock(auth_msg.password)  # This will throw if wrong password
+                    logger.info("Password verified (store already unlocked)")
+                else:
+                    secrets_store.unlock(auth_msg.password)
+                    logger.info("Secrets store unlocked successfully")
+
+                # Check if user wants to change password
+                password_changed = False
+                if auth_msg.change_to_password:
+                    # Validate password change request
+                    if not auth_msg.confirm_new_password:
+                        await send_response(AuthResponseMessage(
+                            success=False,
+                            message="New password confirmation required"
+                        ))
+                        await websocket.close(code=1008, reason="Password confirmation required")
+                        return
+
+                    if auth_msg.change_to_password != auth_msg.confirm_new_password:
+                        await send_response(AuthResponseMessage(
+                            success=False,
+                            message="New passwords do not match"
+                        ))
+                        await websocket.close(code=1008, reason="Password confirmation mismatch")
+                        return
+
+                    # Change the password
+                    try:
+                        secrets_store.change_master_password(auth_msg.password, auth_msg.change_to_password)
+                        password_changed = True
+                        logger.info("Master password changed successfully")
+                    except Exception as e:
+                        logger.error(f"Failed to change password: {e}")
+                        await send_response(AuthResponseMessage(
+                            success=False,
+                            message=f"Failed to change password: {str(e)}"
+                        ))
+                        await websocket.close(code=1011, reason="Password change failed")
+                        return
+
+                is_authenticated = True
+                response_message = "Password changed successfully" if password_changed else "Authentication successful"
+                await send_response(AuthResponseMessage(
+                    success=True,
+                    password_changed=password_changed,
+                    message=response_message
+                ))
+            except InvalidMasterPassword:
+                logger.warning("Invalid password attempt")
+                await send_response(AuthResponseMessage(
+                    success=False,
+                    message="Invalid password"
+                ))
+                await websocket.close(code=1008, reason="Invalid password")
+                return
+            except Exception as e:
+                logger.error(f"Authentication error: {e}")
+                await send_response(AuthResponseMessage(
+                    success=False,
+                    message="Authentication failed"
+                ))
+                await websocket.close(code=1011, reason="Authentication error")
+                return
+
+    except asyncio.TimeoutError:
+        logger.warning("Authentication timeout")
+        await websocket.close(code=1008, reason="Authentication timeout")
+        return
+    except WebSocketDisconnect:
+        logger.info("Client disconnected during authentication")
+        return
+    except Exception as e:
+        logger.error(f"Error during authentication: {e}")
+        await websocket.close(code=1011, reason="Authentication error")
+        return
+
+    # Now authenticated - proceed with normal WebSocket handling
+    registry.websocket = websocket
+
+    # Create WebSocket channel for agent communication
+    channel_id = f"ws_{id(websocket)}"
+    client_id = f"client_{id(websocket)}"
+    logger.info(f"WebSocket authenticated - channel_id: {channel_id}, client_id: {client_id}")
+    ws_channel = WebSocketChannel(channel_id, websocket, session_id="default")
+    gateway.register_channel(ws_channel)
+
     try:
         while True:
             data = await websocket.receive_text()

backend/src/sync/protocol.py

@@ -1,8 +1,24 @@
-from typing import Any, Dict, List, Literal, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from pydantic import BaseModel
 
 
+class AuthMessage(BaseModel):
+    """Authentication message (must be first message from client)"""
+    type: Literal["auth"] = "auth"
+    password: str
+    confirm_password: Optional[str] = None  # Required only for initialization
+    change_to_password: Optional[str] = None  # If provided, change password after auth
+    confirm_new_password: Optional[str] = None  # Required if change_to_password is set
+
+class AuthResponseMessage(BaseModel):
+    """Authentication response from server"""
+    type: Literal["auth_response"] = "auth_response"
+    success: bool
+    needs_confirmation: bool = False  # True if this is first-time setup
+    password_changed: bool = False  # True if password was changed
+    message: str
+
 class SnapshotMessage(BaseModel):
     type: Literal["snapshot"] = "snapshot"
     store: str
@@ -20,7 +36,7 @@ class HelloMessage(BaseModel):
     seqs: Dict[str, int]
 
 # Union type for all messages from backend to frontend
-BackendMessage = Union[SnapshotMessage, PatchMessage]
+BackendMessage = Union[SnapshotMessage, PatchMessage, AuthResponseMessage]
 
 # Union type for all messages from frontend to backend
-FrontendMessage = Union[HelloMessage, PatchMessage]
+FrontendMessage = Union[AuthMessage, HelloMessage, PatchMessage]