Update storage_handler.py

storage_handler.py

@@ -1,215 +1,523 @@
-# storage_handler.py
 import os
 import base64
+import hashlib
+import uuid
 import logging
-from typing import Optional, BinaryIO, Dict, Any
+from typing import Optional, BinaryIO, Dict, Any, List
+from contextlib import contextmanager
+import time
 
 from huggingface_hub import HfFileSystem, HfApi
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 
-logger = logging.getLogger("StorageHandler")
+logger = logging.getLogger("InfinityChat.Storage")
 
 # ------------------------------------------------------------------------
-# Configuration – all from environment variables
+# Configuration
 # ------------------------------------------------------------------------
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if not HF_TOKEN:
-    raise ValueError("HF_TOKEN environment variable is required")
+    raise ValueError("❌ HF_TOKEN environment variable is required")
 
-BUCKET_ID = os.environ.get("INFINITY_CHAT_BUCKET", "infinitychat-data")
+BUCKET_NAME = os.environ.get("INFINITY_CHAT_BUCKET", "infinitychat-data")
 FILE_ENCRYPTION_KEY_B64 = os.environ.get("FILE_ENCRYPTION_KEY", None)
+
 if FILE_ENCRYPTION_KEY_B64 is None:
     from cryptography.fernet import Fernet
     FILE_ENCRYPTION_KEY_B64 = Fernet.generate_key().decode()
-    logger.warning("FILE_ENCRYPTION_KEY not set – generated random key. Set it permanently!")
-FILE_ENCRYPTION_KEY = base64.urlsafe_b64decode(FILE_ENCRYPTION_KEY_B64.encode())
-assert len(FILE_ENCRYPTION_KEY) == 32, "AES‑256 requires exactly 32 bytes"
+    logger.warning("⚠️  FILE_ENCRYPTION_KEY not set - generated random key. Set it permanently!")
+    print(f"Generated FILE_ENCRYPTION_KEY: {FILE_ENCRYPTION_KEY_B64}")
+
+try:
+    FILE_ENCRYPTION_KEY = base64.urlsafe_b64decode(FILE_ENCRYPTION_KEY_B64.encode())
+except Exception as e:
+    raise ValueError(f"Invalid FILE_ENCRYPTION_KEY format: {e}")
+
+assert len(FILE_ENCRYPTION_KEY) == 32, "FILE_ENCRYPTION_KEY must decode to exactly 32 bytes for AES-256"
 
-# Get the bucket owner (your HF username)
+# ------------------------------------------------------------------------
+# Get bucket owner (HF username)
+# ------------------------------------------------------------------------
 _api = HfApi(token=HF_TOKEN)
 try:
-    OWNER = _api.whoami()["name"]
-except Exception:
-    # Fallback: use the bucket ID prefix (if it contains a slash)
-    if "/" in BUCKET_ID:
-        OWNER, BUCKET_ID = BUCKET_ID.split("/", 1)
-    else:
-        raise ValueError("Cannot determine bucket owner. Set INFINITY_CHAT_BUCKET as 'owner/bucket-name'")
+    owner_info = _api.whoami()
+    OWNER = owner_info["name"]
+except Exception as e:
+    logger.warning(f"Could not determine HF username: {e}")
+    # Try to extract from token or use a default
+    try:
+        # Token might be hf_xxxx format, extract first part
+        token_parts = HF_TOKEN.split("_")
+        if len(token_parts) >= 2:
+            OWNER = token_parts[1][:10]  # Use part of token as identifier
+        else:
+            OWNER = "infinitychat"
+    except:
+        OWNER = "infinitychat"
+    
+    logger.info(f"Using owner: {OWNER}")
 
-# Build the proper bucket URI prefix (Xet-native)
-BUCKET_URI_PREFIX = f"hf://buckets/{OWNER}/{BUCKET_ID}"
+# Handle bucket name with or without owner prefix
+if "/" in BUCKET_NAME:
+    OWNER, BUCKET_ID = BUCKET_NAME.split("/", 1)
+else:
+    BUCKET_ID = BUCKET_NAME
+
+BUCKET_URI = f"hf://buckets/{OWNER}/{BUCKET_ID}"
+logger.info(f"📦 Bucket URI: {BUCKET_URI}")
 
 # ------------------------------------------------------------------------
-# HfFileSystem for low-level bucket operations (Xet-native, NO Git/LFS)
+# HfFileSystem instance
 # ------------------------------------------------------------------------
 _fs = HfFileSystem(token=HF_TOKEN)
 
 # ------------------------------------------------------------------------
-# Bucket initialization (using the modern Hugging Face Hub API)
+# Bucket Initialization
 # ------------------------------------------------------------------------
 def ensure_bucket():
-    """
-    Ensure the bucket exists using the native HF Bucket API.
-    The HfApi.create_bucket call does NOT use 'name' – it uses 'bucket_id'.
-    """
-    api = HfApi(token=HF_TOKEN)
-    
-    # Modern API: create_bucket(bucket_id=..., private=True, exist_ok=True)
+    """Create the private bucket if it doesn't exist."""
     try:
-        api.create_bucket(
+        # Try to create bucket
+        _api.create_bucket(
             bucket_id=f"{OWNER}/{BUCKET_ID}",
-            private=False,
+            private=True,
             exist_ok=True
         )
-        logger.info(f"✅ Bucket '{OWNER}/{BUCKET_ID}' is ready (private).")
+        logger.info(f"✅ Bucket '{OWNER}/{BUCKET_ID}' is ready (private)")
     except Exception as e:
-        # Some versions require the full ID
-        if "already exists" in str(e).lower():
-            logger.info(f"Bucket '{OWNER}/{BUCKET_ID}' already exists.")
-        # Try without the owner prefix
+        error_str = str(e).lower()
+        if "already exists" in error_str:
+            logger.info(f"📦 Bucket '{OWNER}/{BUCKET_ID}' already exists")
+        elif "http" in error_str and ("403" in error_str or "401" in error_str):
+            logger.warning(f"⚠️  Cannot create bucket - permission issue: {e}")
+            logger.info("The bucket may already exist or you don't have permissions")
         else:
+            # Try alternative bucket name format
             try:
-                api.create_bucket(
+                _api.create_bucket(
                     bucket_id=BUCKET_ID,
                     private=True,
                     exist_ok=True
                 )
-                logger.info(f"✅ Bucket '{BUCKET_ID}' ready.")
-                global BUCKET_URI_PREFIX
-                BUCKET_URI_PREFIX = f"hf://buckets/{BUCKET_ID}"
+                global OWNER, BUCKET_URI, BUCKET_ID
+                OWNER = ""  # No owner prefix
+                BUCKET_URI = f"hf://buckets/{BUCKET_ID}"
+                logger.info(f"✅ Bucket '{BUCKET_ID}' is ready (alternative format)")
             except Exception as e2:
-                if "already exists" in str(e2).lower():
-                    logger.info(f"Bucket '{BUCKET_ID}' already exists.")
+                error_str2 = str(e2).lower()
+                if "already exists" in error_str2:
+                    logger.info(f"📦 Bucket '{BUCKET_ID}' already exists")
                 else:
-                    logger.error(f"Failed to create bucket: {e2}")
-                    raise
+                    logger.error(f"❌ Failed to create bucket: {e2}")
+                    logger.info("The app will still work but file uploads will fail")
 
-# Initialize bucket
+# Initialize on import
 ensure_bucket()
 
 # ------------------------------------------------------------------------
 # Encryption / Decryption
 # ------------------------------------------------------------------------
-def encrypt_file(data: bytes) -> bytes:
-    """Encrypt binary data using AES‑256‑GCM. Output: nonce (12 B) || ciphertext."""
+def encrypt_bytes(data: bytes, aad: Optional[bytes] = None) -> bytes:
+    """
+    Encrypt binary data with AES-256-GCM.
+    
+    Args:
+        data: Plaintext bytes to encrypt
+        aad: Additional authenticated data (optional)
+    
+    Returns:
+        nonce (12 bytes) + ciphertext + tag (16 bytes) = 12 + len(data) + 16
+    """
     aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
     nonce = os.urandom(12)
-    ciphertext = aesgcm.encrypt(nonce, data, None)
+    ciphertext = aesgcm.encrypt(nonce, data, aad or b"")
     return nonce + ciphertext
 
-def decrypt_file(encrypted_blob: bytes) -> bytes:
-    """Decrypt data produced by encrypt_file()."""
+def decrypt_bytes(encrypted_blob: bytes, aad: Optional[bytes] = None) -> bytes:
+    """
+    Decrypt data encrypted with encrypt_bytes().
+    
+    Args:
+        encrypted_blob: nonce (12 bytes) + ciphertext + tag (16 bytes)
+        aad: Additional authenticated data used during encryption
+    
+    Returns:
+        Decrypted plaintext bytes
+    """
+    aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
     nonce = encrypted_blob[:12]
     ciphertext = encrypted_blob[12:]
-    aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
-    return aesgcm.decrypt(nonce, ciphertext, None)
+    return aesgcm.decrypt(nonce, ciphertext, aad or b"")
+
+def verify_file_integrity(encrypted_blob: bytes) -> bool:
+    """
+    Verify that encrypted blob has valid format.
+    
+    Returns:
+        True if format appears valid (nonce + ciphertext + tag)
+    """
+    if len(encrypted_blob) < 29:  # 12 (nonce) + 1 (minimum data) + 16 (tag)
+        return False
+    return True
 
 # ------------------------------------------------------------------------
-# Core Bucket Operations (Xet-native via HfFileSystem)
+# Path Utilities
 # ------------------------------------------------------------------------
+def _build_path(*parts: str) -> str:
+    """Build a clean path by joining parts."""
+    return "/".join(p.strip("/") for p in parts if p)
+
 def _bucket_path(remote_path: str) -> str:
     """Construct the full hf://buckets/... URI."""
-    return f"{BUCKET_URI_PREFIX}/{remote_path}"
+    return f"{BUCKET_URI}/{remote_path.lstrip('/')}"
 
-def store_file(remote_path: str, data: bytes) -> str:
+def _validate_path(remote_path: str) -> None:
     """
-    Encrypt data and write it to the bucket using HfFileSystem.
-    Returns: the remote path (for DB storage).
+    Validate that a path is safe to use.
+    
+    Raises:
+        ValueError: If path contains dangerous patterns
     """
-    encrypted_data = encrypt_file(data)
-    full_uri = _bucket_path(remote_path)
+    if not remote_path:
+        raise ValueError("Path cannot be empty")
+    if ".." in remote_path.split("/"):
+        raise ValueError("Path traversal detected")
+    if remote_path.startswith("/") or remote_path.startswith("\\"):
+        raise ValueError("Path cannot start with separator")
+    if len(remote_path) > 1024:
+        raise ValueError("Path too long (max 1024 characters)")
+
+# ------------------------------------------------------------------------
+# Core File Operations
+# ------------------------------------------------------------------------
+def store_file(remote_path: str, data: bytes, encrypt: bool = True) -> str:
+    """
+    Store a file in the bucket with optional encryption.
+    
+    Args:
+        remote_path: Path within bucket (e.g., "avatars/user123.jpg")
+        data: File content as bytes
+        encrypt: Whether to encrypt data before storing (default: True)
     
-    with _fs.open(full_uri, "wb") as f:
-        f.write(encrypted_data)
+    Returns:
+        The remote path for later retrieval
     
-    logger.debug(f"Stored file: {full_uri}")
-    return remote_path
+    Raises:
+        ValueError: If path is invalid
+        IOError: If storage fails
+    """
+    _validate_path(remote_path)
+    
+    try:
+        # Encrypt if requested
+        if encrypt:
+            encrypted_data = encrypt_bytes(data, remote_path.encode('utf-8'))
+        else:
+            encrypted_data = data
+        
+        # Ensure parent directory exists (HfFileSystem handles this)
+        full_uri = _bucket_path(remote_path)
+        
+        # Write to bucket
+        with _fs.open(full_uri, "wb") as f:
+            f.write(encrypted_data)
+        
+        # Verify file was written
+        if not _fs.exists(full_uri):
+            raise IOError(f"Failed to verify file was stored: {full_uri}")
+        
+        logger.debug(f"💾 Stored file: {remote_path} ({len(data)} bytes)")
+        return remote_path
+        
+    except Exception as e:
+        logger.error(f"❌ Failed to store file {remote_path}: {e}")
+        raise IOError(f"Storage failed: {e}")
 
-def retrieve_file(remote_path: str) -> bytes:
+def retrieve_file(remote_path: str, decrypt: bool = True) -> bytes:
     """
-    Read from bucket using HfFileSystem and decrypt.
+    Retrieve and optionally decrypt a file from the bucket.
+    
+    Args:
+        remote_path: Path within bucket
+        decrypt: Whether to decrypt data (default: True)
+    
+    Returns:
+        File content as bytes
+    
+    Raises:
+        FileNotFoundError: If file doesn't exist
+        IOError: If retrieval fails
     """
-    full_uri = _bucket_path(remote_path)
+    _validate_path(remote_path)
     
-    with _fs.open(full_uri, "rb") as f:
-        encrypted_data = f.read()
+    full_uri = _bucket_path(remote_path)
     
-    return decrypt_file(encrypted_data)
+    try:
+        # Check if file exists
+        if not _fs.exists(full_uri):
+            raise FileNotFoundError(f"File not found: {remote_path}")
+        
+        # Read from bucket
+        with _fs.open(full_uri, "rb") as f:
+            encrypted_data = f.read()
+        
+        if not encrypted_data:
+            raise IOError(f"Empty file: {remote_path}")
+        
+        # Decrypt if requested
+        if decrypt:
+            if not verify_file_integrity(encrypted_data):
+                raise IOError(f"Corrupted file: {remote_path}")
+            
+            decrypted_data = decrypt_bytes(encrypted_data, remote_path.encode('utf-8'))
+            logger.debug(f"📂 Retrieved file: {remote_path} ({len(decrypted_data)} bytes)")
+            return decrypted_data
+        else:
+            return encrypted_data
+            
+    except FileNotFoundError:
+        raise
+    except Exception as e:
+        logger.error(f"❌ Failed to retrieve file {remote_path}: {e}")
+        raise IOError(f"Retrieval failed: {e}")
 
-def delete_file(remote_path: str):
+def delete_file(remote_path: str) -> bool:
     """
-    Delete a file from the bucket using HfFileSystem.
+    Delete a file from the bucket.
+    
+    Args:
+        remote_path: Path within bucket
+    
+    Returns:
+        True if file was deleted, False if it didn't exist
     """
+    _validate_path(remote_path)
+    
     full_uri = _bucket_path(remote_path)
+    
     try:
-        _fs.rm(full_uri)
-        logger.debug(f"Deleted file: {full_uri}")
-    except FileNotFoundError:
-        logger.warning(f"File not found for deletion: {full_uri}")
+        if _fs.exists(full_uri):
+            _fs.rm(full_uri)
+            logger.debug(f"🗑️  Deleted file: {remote_path}")
+            return True
+        else:
+            logger.warning(f"⚠️  File not found for deletion: {remote_path}")
+            return False
     except Exception as e:
-        logger.error(f"Failed to delete file {full_uri}: {e}")
+        logger.error(f"❌ Failed to delete file {remote_path}: {e}")
+        raise IOError(f"Deletion failed: {e}")
 
-def list_files(prefix: str = "") -> list:
+# ------------------------------------------------------------------------
+# Bulk Operations
+# ------------------------------------------------------------------------
+def list_files(prefix: str = "", recursive: bool = True) -> List[str]:
     """
     List files in the bucket under a given prefix.
-    Returns list of full URIs.
+    
+    Args:
+        prefix: Optional path prefix to filter by
+        recursive: Whether to list files recursively
+    
+    Returns:
+        List of file paths (relative to bucket root)
     """
-    search_path = f"{BUCKET_URI_PREFIX}/{prefix}" if prefix else BUCKET_URI_PREFIX
+    search_path = _bucket_path(prefix) if prefix else BUCKET_URI
+    
     try:
-        return _fs.ls(search_path)
+        if recursive:
+            items = _fs.ls(search_path, detail=False, recursive=True)
+        else:
+            items = _fs.ls(search_path, detail=False)
+        
+        # Convert full URIs to relative paths
+        prefix_len = len(BUCKET_URI) + 1  # +1 for trailing slash
+        return [item[prefix_len:] for item in items if not _fs.isdir(item)]
+        
     except FileNotFoundError:
         return []
+    except Exception as e:
+        logger.error(f"❌ Failed to list files: {e}")
+        return []
 
 def file_exists(remote_path: str) -> bool:
     """Check if a file exists in the bucket."""
-    full_uri = _bucket_path(remote_path)
-    return _fs.exists(full_uri)
+    _validate_path(remote_path)
+    return _fs.exists(_bucket_path(remote_path))
 
 def get_file_size(remote_path: str) -> Optional[int]:
-    """Get the size of a file in bytes (compressed)."""
-    full_uri = _bucket_path(remote_path)
+    """
+    Get the size of a file in bytes.
+    
+    Returns:
+        File size in bytes, or None if file doesn't exist
+    """
+    _validate_path(remote_path)
+    
     try:
-        info = _fs.info(full_uri)
+        info = _fs.info(_bucket_path(remote_path))
         return info.get("size")
-    except Exception:
+    except FileNotFoundError:
+        return None
+    except Exception as e:
+        logger.error(f"Failed to get file size: {e}")
+        return None
+
+def get_file_info(remote_path: str) -> Optional[Dict[str, Any]]:
+    """
+    Get detailed information about a file.
+    
+    Returns:
+        Dict with file metadata, or None if file doesn't exist
+    """
+    _validate_path(remote_path)
+    
+    try:
+        info = _fs.info(_bucket_path(remote_path))
+        return {
+            "name": remote_path,
+            "size": info.get("size"),
+            "created": info.get("created"),
+            "modified": info.get("last_modified"),
+            "type": info.get("type", "file")
+        }
+    except FileNotFoundError:
+        return None
+    except Exception as e:
+        logger.error(f"Failed to get file info: {e}")
         return None
 
 # ------------------------------------------------------------------------
-# Stream-based store (for reassembled chunks)
+# Stream-based Operations
 # ------------------------------------------------------------------------
-def store_file_stream(remote_path: str, data_stream: BinaryIO) -> str:
+def store_file_stream(remote_path: str, data_stream: BinaryIO, encrypt: bool = True) -> str:
     """
-    Read the entire stream, encrypt, and write to bucket.
+    Store file from a binary stream.
+    
+    Args:
+        remote_path: Path within bucket
+        data_stream: Binary stream to read from
+        encrypt: Whether to encrypt data
+    
+    Returns:
+        The remote path
     """
     data = data_stream.read()
-    return store_file(remote_path, data)
+    return store_file(remote_path, data, encrypt=encrypt)
+
+def retrieve_file_stream(remote_path: str, decrypt: bool = True) -> BinaryIO:
+    """
+    Retrieve file as a binary stream.
+    
+    Args:
+        remote_path: Path within bucket
+        decrypt: Whether to decrypt data
+    
+    Returns:
+        BinaryIO object with file content
+    """
+    data = retrieve_file(remote_path, decrypt=decrypt)
+    import io
+    return io.BytesIO(data)
 
 # ------------------------------------------------------------------------
-# Bucket info (for monitoring / health checks)
+# Storage Statistics
 # ------------------------------------------------------------------------
-def get_bucket_info() -> Dict[str, Any]:
-    """Return information about the bucket."""
-    api = HfApi(token=HF_TOKEN)
+def get_storage_stats() -> Dict[str, Any]:
+    """
+    Get storage statistics.
+    
+    Returns:
+        Dict with storage information
+    """
     try:
-        # Try the modern API first
-        info = api.bucket_info(bucket_id=f"{OWNER}/{BUCKET_ID}")
+        files = list_files()
+        total_size = 0
+        file_count = 0
+        
+        for f in files:
+            size = get_file_size(f)
+            if size is not None:
+                total_size += size
+                file_count += 1
+        
         return {
-            "id": f"{OWNER}/{BUCKET_ID}",
-            "exists": True,
-            "private": info.get("private", True),
-            "total_files": len(list_files()),
+            "bucket": f"{OWNER}/{BUCKET_ID}",
+            "file_count": file_count,
+            "total_size": total_size,
+            "total_size_mb": round(total_size / (1024 * 1024), 2)
         }
-    except Exception:
+    except Exception as e:
+        logger.error(f"Failed to get storage stats: {e}")
         return {
-            "id": f"{OWNER}/{BUCKET_ID}",
-            "exists": False,
-            "note": "Cannot fetch bucket info"
+            "bucket": f"{OWNER}/{BUCKET_ID}",
+            "file_count": 0,
+            "total_size": 0,
+            "total_size_mb": 0,
+            "error": str(e)
         }
 
 # ------------------------------------------------------------------------
-# Clean shutdown
+# Backup and Maintenance
+# ------------------------------------------------------------------------
+def create_backup(backup_prefix: str = "backups") -> str:
+    """
+    Create a backup of all files in the bucket.
+    
+    Args:
+        backup_prefix: Path prefix for backup location
+    
+    Returns:
+        Backup path
+    """
+    timestamp = int(time.time())
+    backup_path = f"{backup_prefix}/backup_{timestamp}"
+    
+    try:
+        files = list_files()
+        for file_path in files:
+            try:
+                data = retrieve_file(file_path)
+                backup_file_path = f"{backup_path}/{file_path}"
+                store_file(backup_file_path, data, encrypt=True)
+            except Exception as e:
+                logger.error(f"Failed to backup {file_path}: {e}")
+        
+        logger.info(f"💾 Backup created at: {backup_path}")
+        return backup_path
+    except Exception as e:
+        logger.error(f"❌ Backup failed: {e}")
+        raise
+
+# ------------------------------------------------------------------------
+# Cleanup
 # ------------------------------------------------------------------------
 def close():
-    """Clean up resources (HfFileSystem)."""
+    """Clean up resources."""
     if _fs:
-        _fs.close()
+        try:
+            _fs.close()
+        except:
+            pass
+
+# ------------------------------------------------------------------------
+# Export public API
+# ------------------------------------------------------------------------
+__all__ = [
+    'store_file',
+    'retrieve_file', 
+    'delete_file',
+    'list_files',
+    'file_exists',
+    'get_file_size',
+    'get_file_info',
+    'store_file_stream',
+    'retrieve_file_stream',
+    'get_storage_stats',
+    'create_backup',
+    'encrypt_bytes',
+    'decrypt_bytes',
+    'close',
+    'OWNER',
+    'BUCKET_ID',
+    'BUCKET_URI'
+]