Add wallet key encryption service

backend/.env.example

@@ -10,6 +10,7 @@ SECRET_KEY=your-secret-key-change-this-in-production
 ALGORITHM=HS256
 ACCESS_TOKEN_EXPIRE_MINUTES=30
 REFRESH_TOKEN_EXPIRE_DAYS=7
+WALLET_MASTER_KEY=replace-with-strong-unique-master-key
 
 # AWS S3 (for file storage)
 AWS_ACCESS_KEY_ID=your-access-key

backend/app/config.py

@@ -24,6 +24,7 @@ class Settings(BaseSettings):
     ALGORITHM: str = "HS256"
     ACCESS_TOKEN_EXPIRE_MINUTES: int = 30
     REFRESH_TOKEN_EXPIRE_DAYS: int = 7
+    WALLET_MASTER_KEY: str
 
     # CORS
     ALLOWED_ORIGINS: str = "http://localhost:3000,http://localhost:8000"

backend/app/main.py

@@ -11,6 +11,7 @@ from app.config import settings
 from app.database import async_engine, Base
 from app.redis_client import close_redis
 from app.routers import auth, email, social, video, files, blockchain, ai_chat, devices, miner
+from app.services.crypto import rotate_plaintext_wallet_keys
 
 
 @asynccontextmanager
@@ -26,6 +27,13 @@ async def lifespan(app: FastAPI):
     print("Database tables created successfully")
     print(f"Server running on {settings.ENVIRONMENT} mode")
 
+    # Re-encrypt any legacy plaintext wallet keys before serving requests
+    updated_users, updated_wallets = await rotate_plaintext_wallet_keys()
+    if updated_users or updated_wallets:
+        print(
+            f"Re-encrypted {updated_users} user keys and {updated_wallets} wallet keys"
+        )
+
     yield
 
     # Shutdown

backend/app/routers/auth.py

@@ -15,6 +15,7 @@ from app.auth import (
     get_current_active_user
 )
 from app.services.blockchain import BlockchainService
+from app.services.crypto import wallet_crypto, WalletKeyEncryptionError
 from datetime import datetime
 
 router = APIRouter(prefix="/api/auth", tags=["Authentication"])
@@ -40,6 +41,14 @@ async def register(user_data: UserCreate, db: AsyncSession = Depends(get_db)):
     # Generate wallet
     wallet_address, private_key = BlockchainService.generate_wallet_address()
 
+    try:
+        encrypted_private_key = wallet_crypto.encrypt(private_key)
+    except WalletKeyEncryptionError:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Unable to encrypt wallet key"
+        )
+
     # Create user
     user = User(
         username=user_data.username,
@@ -47,7 +56,7 @@ async def register(user_data: UserCreate, db: AsyncSession = Depends(get_db)):
         full_name=user_data.full_name,
         hashed_password=get_password_hash(user_data.password),
         wallet_address=wallet_address,
-        wallet_private_key=private_key,  # In production, encrypt this!
+        wallet_private_key=encrypted_private_key,
         balance=100.0,  # Starting bonus
         created_at=datetime.utcnow()
     )

backend/app/routers/blockchain.py

@@ -11,6 +11,7 @@ from app.models.user import User
 from app.models.blockchain import Block, Transaction, Wallet
 from app.auth import get_current_active_user
 from app.services.blockchain import BlockchainService
+from app.services.crypto import WalletKeyDecryptionError
 
 router = APIRouter(prefix="/api/blockchain", tags=["Blockchain"])
 
@@ -112,13 +113,19 @@ async def create_transaction(
         )
 
     # Create transaction
-    transaction = await BlockchainService.create_transaction(
+    try:
-        db=db,
+        transaction = await BlockchainService.create_transaction(
-        from_address=current_user.wallet_address,
+            db=db,
-        to_address=tx_data.to_address,
+            from_address=current_user.wallet_address,
-        amount=tx_data.amount,
+            to_address=tx_data.to_address,
-        private_key=current_user.wallet_private_key
+            amount=tx_data.amount,
-    )
+            encrypted_private_key=current_user.wallet_private_key
+        )
+    except WalletKeyDecryptionError:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Wallet key could not be decrypted"
+        )
 
     # Update balances (simplified - in production would be done on block confirmation)
     current_user.balance -= tx_data.amount

backend/app/services/blockchain.py

@@ -8,6 +8,7 @@ from sqlalchemy import select, desc
 from app.models.blockchain import Block, Transaction, Wallet
 from app.models.user import User
 from app.config import settings
+from app.services.crypto import wallet_crypto, WalletKeyDecryptionError
 import secrets
 
 
@@ -140,9 +141,16 @@ class BlockchainService:
         from_address: str,
         to_address: str,
         amount: float,
-        private_key: str
+        encrypted_private_key: str
     ) -> Transaction:
         """Create a new transaction"""
+        try:
+            private_key = wallet_crypto.decrypt(encrypted_private_key)
+        except WalletKeyDecryptionError as exc:
+            raise WalletKeyDecryptionError(
+                "Unable to decrypt wallet key for transaction"
+            ) from exc
+
         # Generate transaction hash
         tx_data = f"{from_address}{to_address}{amount}{datetime.utcnow()}"
         transaction_hash = hashlib.sha256(tx_data.encode()).hexdigest()

backend/app/services/crypto.py

@@ -0,0 +1,106 @@
+"""Utilities for encrypting wallet secrets."""
+from __future__ import annotations
+
+import base64
+import hashlib
+import string
+from dataclasses import dataclass
+from typing import Tuple
+
+from cryptography.fernet import Fernet, InvalidToken
+from sqlalchemy import select
+
+from app.config import settings
+from app.database import AsyncSessionLocal
+from app.models.blockchain import Wallet
+from app.models.user import User
+
+
+class WalletKeyEncryptionError(RuntimeError):
+    """Raised when a wallet key cannot be encrypted."""
+
+
+class WalletKeyDecryptionError(RuntimeError):
+    """Raised when a wallet key cannot be decrypted."""
+
+
+@dataclass
+class WalletCryptoService:
+    """Simple Fernet wrapper used for wallet key protection."""
+
+    master_key: str
+
+    def __post_init__(self) -> None:
+        if not self.master_key or len(self.master_key) < 32:
+            raise ValueError("WALLET_MASTER_KEY must be at least 32 characters long")
+        derived_key = self._derive_key(self.master_key)
+        self._fernet = Fernet(derived_key)
+
+    @staticmethod
+    def _derive_key(master_key: str) -> bytes:
+        """Derive a url-safe 32-byte key for Fernet from the provided secret."""
+        digest = hashlib.sha256(master_key.encode("utf-8")).digest()
+        return base64.urlsafe_b64encode(digest)
+
+    def encrypt(self, plaintext: str) -> str:
+        """Encrypt a wallet secret."""
+        if plaintext is None:
+            raise WalletKeyEncryptionError("Cannot encrypt an empty wallet key")
+        try:
+            token = self._fernet.encrypt(plaintext.encode("utf-8"))
+            return token.decode("utf-8")
+        except Exception as exc:  # pragma: no cover - defensive
+            raise WalletKeyEncryptionError("Unable to encrypt wallet key") from exc
+
+    def decrypt(self, token: str) -> str:
+        """Decrypt a wallet secret."""
+        if token is None:
+            raise WalletKeyDecryptionError("Wallet key is missing")
+        try:
+            plaintext = self._fernet.decrypt(token.encode("utf-8"))
+            return plaintext.decode("utf-8")
+        except (InvalidToken, ValueError, TypeError) as exc:
+            raise WalletKeyDecryptionError("Unable to decrypt wallet key") from exc
+
+
+wallet_crypto = WalletCryptoService(settings.WALLET_MASTER_KEY)
+
+
+def _looks_like_plaintext_key(value: str | None) -> bool:
+    if not value:
+        return False
+    normalized = value.strip()
+    return len(normalized) == 64 and all(ch in string.hexdigits for ch in normalized)
+
+
+async def rotate_plaintext_wallet_keys() -> Tuple[int, int]:
+    """Encrypt plaintext wallet keys that may still exist in the database."""
+    updated_users = 0
+    updated_wallets = 0
+
+    async with AsyncSessionLocal() as session:
+        user_result = await session.execute(select(User).where(User.wallet_private_key.is_not(None)))
+        for user in user_result.scalars():
+            if _looks_like_plaintext_key(user.wallet_private_key):
+                user.wallet_private_key = wallet_crypto.encrypt(user.wallet_private_key)
+                updated_users += 1
+
+        wallet_result = await session.execute(select(Wallet).where(Wallet.private_key.is_not(None)))
+        for wallet in wallet_result.scalars():
+            if _looks_like_plaintext_key(wallet.private_key):
+                wallet.private_key = wallet_crypto.encrypt(wallet.private_key)
+                updated_wallets += 1
+
+        if updated_users or updated_wallets:
+            await session.commit()
+
+    return updated_users, updated_wallets
+
+
+__all__ = [
+    "WalletCryptoService",
+    "WalletKeyEncryptionError",
+    "WalletKeyDecryptionError",
+    "wallet_crypto",
+    "rotate_plaintext_wallet_keys",
+]