quantumflow-sdk 0.2.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

api/routes/chat_routes.py

@@ -0,0 +1,565 @@
+"""
+QChat API Routes - Quantum-Secure Messaging
+
+P2P encrypted messaging using QuantumFlow as relay.
+Backend only sees encrypted blobs - cannot decrypt.
+"""
+
+from fastapi import APIRouter, HTTPException, Depends, WebSocket, WebSocketDisconnect
+from pydantic import BaseModel, Field
+from typing import Optional, List
+from datetime import datetime, timedelta
+import uuid
+import hashlib
+import asyncio
+from collections import defaultdict
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select, or_, and_
+from sqlalchemy.orm import selectinload
+
+from db.database import get_async_db
+from db.models import ChatUser, QuantumChannel, ChatMessage, ChannelStatus, MessageStatus
+
+router = APIRouter(prefix="/chat", tags=["QChat"])
+
+# ==================== WebSocket Connection Manager ====================
+
+class ConnectionManager:
+    """Manage WebSocket connections for real-time messaging."""
+
+    def __init__(self):
+        # user_id -> WebSocket
+        self.active_connections: dict[str, WebSocket] = {}
+        # user_id -> list of pending messages
+        self.pending_messages: dict[str, list] = defaultdict(list)
+
+    async def connect(self, user_id: str, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections[user_id] = websocket
+
+        # Send any pending messages
+        if user_id in self.pending_messages:
+            for msg in self.pending_messages[user_id]:
+                await websocket.send_json(msg)
+            del self.pending_messages[user_id]
+
+    def disconnect(self, user_id: str):
+        if user_id in self.active_connections:
+            del self.active_connections[user_id]
+
+    async def send_to_user(self, user_id: str, message: dict):
+        if user_id in self.active_connections:
+            await self.active_connections[user_id].send_json(message)
+            return True
+        else:
+            # Store for later delivery
+            self.pending_messages[user_id].append(message)
+            return False
+
+    def is_online(self, user_id: str) -> bool:
+        return user_id in self.active_connections
+
+
+manager = ConnectionManager()
+
+
+# ==================== Request/Response Models ====================
+
+class RegisterUserRequest(BaseModel):
+    phone_number: str = Field(..., description="Phone number in E.164 format")
+    display_name: Optional[str] = None
+    firebase_uid: Optional[str] = None
+    device_token: Optional[str] = None
+    platform: Optional[str] = None
+
+
+class RegisterUserResponse(BaseModel):
+    user_id: str
+    phone_number: str
+    display_name: Optional[str]
+    created_at: datetime
+
+
+class EstablishChannelRequest(BaseModel):
+    recipient_phone: str = Field(..., description="Recipient phone number")
+    sender_id: str = Field(..., description="Sender user ID")
+
+
+class ChannelResponse(BaseModel):
+    channel_id: str
+    status: str
+    bell_pairs_remaining: int
+    error_rate: float
+    created_at: datetime
+    established_at: Optional[datetime]
+
+
+class SendMessageRequest(BaseModel):
+    channel_id: str
+    sender_id: str
+    encrypted_content: str = Field(..., description="Encrypted message (base64)")
+    content_hash: str = Field(..., description="SHA-256 hash of plaintext")
+    compression_ratio: Optional[float] = None
+    bell_pair_id: Optional[str] = None
+    teleport_fidelity: Optional[float] = None
+
+
+class MessageResponse(BaseModel):
+    message_id: str
+    channel_id: str
+    sender_id: str
+    encrypted_content: str
+    status: str
+    compression_ratio: Optional[float]
+    eavesdrop_detected: bool
+    created_at: datetime
+
+
+class GetMessagesRequest(BaseModel):
+    channel_id: str
+    user_id: str
+    limit: int = 50
+    before_id: Optional[str] = None
+
+
+# ==================== User Registration ====================
+
+@router.post("/register", response_model=RegisterUserResponse)
+async def register_user(
+    request: RegisterUserRequest,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """Register a new QChat user by phone number."""
+
+    # Check if user exists
+    result = await db.execute(
+        select(ChatUser).where(ChatUser.phone_number == request.phone_number)
+    )
+    existing = result.scalar_one_or_none()
+
+    if existing:
+        # Update existing user
+        existing.display_name = request.display_name or existing.display_name
+        existing.firebase_uid = request.firebase_uid or existing.firebase_uid
+        existing.device_token = request.device_token or existing.device_token
+        existing.platform = request.platform or existing.platform
+        existing.updated_at = datetime.utcnow()
+        await db.commit()
+
+        return RegisterUserResponse(
+            user_id=str(existing.id),
+            phone_number=existing.phone_number,
+            display_name=existing.display_name,
+            created_at=existing.created_at
+        )
+
+    # Create new user
+    user = ChatUser(
+        phone_number=request.phone_number,
+        display_name=request.display_name,
+        firebase_uid=request.firebase_uid,
+        device_token=request.device_token,
+        platform=request.platform
+    )
+    db.add(user)
+    await db.commit()
+    await db.refresh(user)
+
+    return RegisterUserResponse(
+        user_id=str(user.id),
+        phone_number=user.phone_number,
+        display_name=user.display_name,
+        created_at=user.created_at
+    )
+
+
+@router.get("/user/{phone_number}")
+async def get_user_by_phone(
+    phone_number: str,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """Look up user by phone number."""
+
+    result = await db.execute(
+        select(ChatUser).where(ChatUser.phone_number == phone_number)
+    )
+    user = result.scalar_one_or_none()
+
+    if not user:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    return {
+        "user_id": str(user.id),
+        "phone_number": user.phone_number,
+        "display_name": user.display_name,
+        "is_online": manager.is_online(str(user.id))
+    }
+
+
+# ==================== Channel Management ====================
+
+@router.post("/channel/establish", response_model=ChannelResponse)
+async def establish_channel(
+    request: EstablishChannelRequest,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """Establish a quantum-secure channel with another user."""
+
+    # Get sender
+    sender_result = await db.execute(
+        select(ChatUser).where(ChatUser.id == uuid.UUID(request.sender_id))
+    )
+    sender = sender_result.scalar_one_or_none()
+    if not sender:
+        raise HTTPException(status_code=404, detail="Sender not found")
+
+    # Get recipient by phone
+    recipient_result = await db.execute(
+        select(ChatUser).where(ChatUser.phone_number == request.recipient_phone)
+    )
+    recipient = recipient_result.scalar_one_or_none()
+    if not recipient:
+        raise HTTPException(status_code=404, detail="Recipient not found")
+
+    # Check for existing channel
+    channel_result = await db.execute(
+        select(QuantumChannel).where(
+            or_(
+                and_(
+                    QuantumChannel.user_a_id == sender.id,
+                    QuantumChannel.user_b_id == recipient.id
+                ),
+                and_(
+                    QuantumChannel.user_a_id == recipient.id,
+                    QuantumChannel.user_b_id == sender.id
+                )
+            )
+        )
+    )
+    existing_channel = channel_result.scalar_one_or_none()
+
+    if existing_channel:
+        return ChannelResponse(
+            channel_id=str(existing_channel.id),
+            status=existing_channel.status.value,
+            bell_pairs_remaining=existing_channel.bell_pairs_remaining,
+            error_rate=existing_channel.error_rate,
+            created_at=existing_channel.created_at,
+            established_at=existing_channel.established_at
+        )
+
+    # Create new channel
+    channel = QuantumChannel(
+        user_a_id=sender.id,
+        user_b_id=recipient.id,
+        status=ChannelStatus.ESTABLISHING,
+        qkd_key_id=str(uuid.uuid4()),
+        bell_pairs_remaining=1000,
+        key_generated_at=datetime.utcnow(),
+        key_expires_at=datetime.utcnow() + timedelta(days=7)
+    )
+    db.add(channel)
+
+    # Simulate QKD exchange (in production, call actual QKD API)
+    channel.status = ChannelStatus.READY
+    channel.established_at = datetime.utcnow()
+
+    await db.commit()
+    await db.refresh(channel)
+
+    # Notify recipient
+    await manager.send_to_user(str(recipient.id), {
+        "type": "channel_established",
+        "channel_id": str(channel.id),
+        "from_user": sender.display_name or sender.phone_number
+    })
+
+    return ChannelResponse(
+        channel_id=str(channel.id),
+        status=channel.status.value,
+        bell_pairs_remaining=channel.bell_pairs_remaining,
+        error_rate=channel.error_rate,
+        created_at=channel.created_at,
+        established_at=channel.established_at
+    )
+
+
+@router.get("/channels/{user_id}")
+async def get_user_channels(
+    user_id: str,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """Get all channels for a user."""
+
+    uid = uuid.UUID(user_id)
+    result = await db.execute(
+        select(QuantumChannel)
+        .where(
+            or_(
+                QuantumChannel.user_a_id == uid,
+                QuantumChannel.user_b_id == uid
+            )
+        )
+        .options(selectinload(QuantumChannel.user_a), selectinload(QuantumChannel.user_b))
+    )
+    channels = result.scalars().all()
+
+    return [
+        {
+            "channel_id": str(ch.id),
+            "status": ch.status.value,
+            "bell_pairs_remaining": ch.bell_pairs_remaining,
+            "other_user": {
+                "id": str(ch.user_b.id if str(ch.user_a_id) == user_id else ch.user_a.id),
+                "phone": ch.user_b.phone_number if str(ch.user_a_id) == user_id else ch.user_a.phone_number,
+                "name": ch.user_b.display_name if str(ch.user_a_id) == user_id else ch.user_a.display_name
+            },
+            "last_message_at": ch.last_message_at
+        }
+        for ch in channels
+    ]
+
+
+# ==================== Messaging ====================
+
+@router.post("/message/send", response_model=MessageResponse)
+async def send_message(
+    request: SendMessageRequest,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """Send an encrypted message through a quantum channel."""
+
+    # Get channel
+    channel_result = await db.execute(
+        select(QuantumChannel).where(QuantumChannel.id == uuid.UUID(request.channel_id))
+    )
+    channel = channel_result.scalar_one_or_none()
+    if not channel:
+        raise HTTPException(status_code=404, detail="Channel not found")
+
+    if channel.status != ChannelStatus.READY:
+        raise HTTPException(status_code=400, detail=f"Channel not ready: {channel.status.value}")
+
+    # Verify sender is in channel
+    sender_uuid = uuid.UUID(request.sender_id)
+    if sender_uuid not in [channel.user_a_id, channel.user_b_id]:
+        raise HTTPException(status_code=403, detail="Sender not in channel")
+
+    # Determine recipient
+    recipient_id = channel.user_b_id if sender_uuid == channel.user_a_id else channel.user_a_id
+
+    # Consume bell pair
+    if channel.bell_pairs_remaining > 0:
+        channel.bell_pairs_remaining -= 1
+    else:
+        raise HTTPException(status_code=400, detail="No bell pairs remaining. Regenerate keys.")
+
+    # Simulate eavesdrop detection (random ~2% chance for demo)
+    import random
+    eavesdrop_detected = random.random() < 0.02
+
+    if eavesdrop_detected:
+        channel.eavesdrop_detected_count += 1
+        channel.last_eavesdrop_at = datetime.utcnow()
+        channel.error_rate = min(0.25, channel.error_rate + 0.05)
+
+        if channel.error_rate >= 0.25:
+            channel.status = ChannelStatus.COMPROMISED
+
+    # Create message
+    message = ChatMessage(
+        channel_id=channel.id,
+        sender_id=sender_uuid,
+        encrypted_content=request.encrypted_content,
+        content_hash=request.content_hash,
+        compression_ratio=request.compression_ratio,
+        bell_pair_id=request.bell_pair_id,
+        teleport_fidelity=request.teleport_fidelity,
+        eavesdrop_detected=eavesdrop_detected,
+        status=MessageStatus.SENT,
+        sent_at=datetime.utcnow()
+    )
+    db.add(message)
+
+    channel.last_message_at = datetime.utcnow()
+
+    await db.commit()
+    await db.refresh(message)
+
+    # Send to recipient via WebSocket
+    delivered = await manager.send_to_user(str(recipient_id), {
+        "type": "new_message",
+        "message_id": str(message.id),
+        "channel_id": str(channel.id),
+        "sender_id": request.sender_id,
+        "encrypted_content": request.encrypted_content,
+        "content_hash": request.content_hash,
+        "compression_ratio": request.compression_ratio,
+        "eavesdrop_detected": eavesdrop_detected,
+        "created_at": message.created_at.isoformat()
+    })
+
+    if delivered:
+        message.status = MessageStatus.DELIVERED
+        message.delivered_at = datetime.utcnow()
+        await db.commit()
+
+    return MessageResponse(
+        message_id=str(message.id),
+        channel_id=str(channel.id),
+        sender_id=str(message.sender_id),
+        encrypted_content=message.encrypted_content,
+        status=message.status.value,
+        compression_ratio=message.compression_ratio,
+        eavesdrop_detected=message.eavesdrop_detected,
+        created_at=message.created_at
+    )
+
+
+@router.get("/messages/{channel_id}")
+async def get_messages(
+    channel_id: str,
+    user_id: str,
+    limit: int = 50,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """Get messages from a channel."""
+
+    # Verify user is in channel
+    channel_result = await db.execute(
+        select(QuantumChannel).where(QuantumChannel.id == uuid.UUID(channel_id))
+    )
+    channel = channel_result.scalar_one_or_none()
+    if not channel:
+        raise HTTPException(status_code=404, detail="Channel not found")
+
+    user_uuid = uuid.UUID(user_id)
+    if user_uuid not in [channel.user_a_id, channel.user_b_id]:
+        raise HTTPException(status_code=403, detail="User not in channel")
+
+    # Get messages
+    result = await db.execute(
+        select(ChatMessage)
+        .where(ChatMessage.channel_id == uuid.UUID(channel_id))
+        .order_by(ChatMessage.created_at.desc())
+        .limit(limit)
+    )
+    messages = result.scalars().all()
+
+    # Mark as read
+    for msg in messages:
+        if msg.sender_id != user_uuid and msg.status != MessageStatus.READ:
+            msg.status = MessageStatus.READ
+            msg.read_at = datetime.utcnow()
+    await db.commit()
+
+    return [
+        {
+            "message_id": str(msg.id),
+            "sender_id": str(msg.sender_id),
+            "encrypted_content": msg.encrypted_content,
+            "content_hash": msg.content_hash,
+            "compression_ratio": msg.compression_ratio,
+            "eavesdrop_detected": msg.eavesdrop_detected,
+            "status": msg.status.value,
+            "created_at": msg.created_at.isoformat()
+        }
+        for msg in reversed(messages)
+    ]
+
+
+@router.post("/channel/{channel_id}/regenerate-keys")
+async def regenerate_keys(
+    channel_id: str,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """Regenerate QKD keys for a compromised channel."""
+
+    result = await db.execute(
+        select(QuantumChannel).where(QuantumChannel.id == uuid.UUID(channel_id))
+    )
+    channel = result.scalar_one_or_none()
+    if not channel:
+        raise HTTPException(status_code=404, detail="Channel not found")
+
+    # Regenerate keys
+    channel.qkd_key_id = str(uuid.uuid4())
+    channel.bell_pairs_remaining = 1000
+    channel.key_generated_at = datetime.utcnow()
+    channel.key_expires_at = datetime.utcnow() + timedelta(days=7)
+    channel.status = ChannelStatus.READY
+    channel.error_rate = 0.0
+
+    await db.commit()
+
+    # Notify both users
+    await manager.send_to_user(str(channel.user_a_id), {
+        "type": "keys_regenerated",
+        "channel_id": channel_id
+    })
+    await manager.send_to_user(str(channel.user_b_id), {
+        "type": "keys_regenerated",
+        "channel_id": channel_id
+    })
+
+    return {
+        "status": "success",
+        "channel_id": channel_id,
+        "new_key_id": channel.qkd_key_id,
+        "bell_pairs": channel.bell_pairs_remaining
+    }
+
+
+# ==================== WebSocket ====================
+
+@router.websocket("/ws/{user_id}")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    user_id: str,
+    db: AsyncSession = Depends(get_async_db)
+):
+    """WebSocket connection for real-time messaging."""
+
+    await manager.connect(user_id, websocket)
+
+    # Update user online status
+    result = await db.execute(
+        select(ChatUser).where(ChatUser.id == uuid.UUID(user_id))
+    )
+    user = result.scalar_one_or_none()
+    if user:
+        user.is_online = True
+        await db.commit()
+
+    try:
+        while True:
+            data = await websocket.receive_json()
+
+            # Handle different message types
+            if data.get("type") == "ping":
+                await websocket.send_json({"type": "pong"})
+
+            elif data.get("type") == "typing":
+                # Forward typing indicator
+                channel_id = data.get("channel_id")
+                result = await db.execute(
+                    select(QuantumChannel).where(QuantumChannel.id == uuid.UUID(channel_id))
+                )
+                channel = result.scalar_one_or_none()
+                if channel:
+                    recipient_id = str(channel.user_b_id) if str(channel.user_a_id) == user_id else str(channel.user_a_id)
+                    await manager.send_to_user(recipient_id, {
+                        "type": "typing",
+                        "channel_id": channel_id,
+                        "user_id": user_id
+                    })
+
+    except WebSocketDisconnect:
+        manager.disconnect(user_id)
+
+        # Update user offline status
+        if user:
+            user.is_online = False
+            user.last_seen = datetime.utcnow()
+            await db.commit()