Update storage_handler.py

storage_handler.py

@@ -4,7 +4,6 @@ import hashlib
 import uuid
 import logging
 from typing import Optional, BinaryIO, Dict, Any, List
-from contextlib import contextmanager
 import time
 
 from huggingface_hub import HfFileSystem, HfApi
@@ -44,20 +43,8 @@ try:
     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}")
+    OWNER = "infinitychat"
 
-# Handle bucket name with or without owner prefix
 if "/" in BUCKET_NAME:
     OWNER, BUCKET_ID = BUCKET_NAME.split("/", 1)
 else:
@@ -99,7 +86,6 @@ def ensure_bucket():
                     private=True,
                     exist_ok=True
                 )
-                # now we can safely reassign globals
                 OWNER = ""
                 BUCKET_URI = f"hf://buckets/{BUCKET_ID}"
                 logger.info(f"✅ Bucket '{BUCKET_ID}' is ready (alternative format)")
@@ -111,73 +97,36 @@ def ensure_bucket():
                     logger.error(f"❌ Failed to create bucket: {e2}")
                     logger.info("The app will still work but file uploads will fail")
 
-# Initialize on import
 ensure_bucket()
 
 # ------------------------------------------------------------------------
 # Encryption / Decryption
 # ------------------------------------------------------------------------
 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, aad or b"")
     return nonce + ciphertext
 
 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:]
     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
+    return len(encrypted_blob) >= 29
 
 # ------------------------------------------------------------------------
 # 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}/{remote_path.lstrip('/')}"
 
 def _validate_path(remote_path: str) -> None:
-    """
-    Validate that a path is safe to use.
-    
-    Raises:
-        ValueError: If path contains dangerous patterns
-    """
     if not remote_path:
         raise ValueError("Path cannot be empty")
     if ".." in remote_path.split("/"):
@@ -191,90 +140,41 @@ def _validate_path(remote_path: str) -> None:
 # 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)
-    
-    Returns:
-        The remote path for later retrieval
-    
-    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, decrypt: bool = True) -> bytes:
-    """
-    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
-    """
     _validate_path(remote_path)
-    
     full_uri = _bucket_path(remote_path)
-    
     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:
@@ -282,19 +182,8 @@ def retrieve_file(remote_path: str, decrypt: bool = True) -> bytes:
         raise IOError(f"Retrieval failed: {e}")
 
 def delete_file(remote_path: str) -> bool:
-    """
-    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:
         if _fs.exists(full_uri):
             _fs.rm(full_uri)
@@ -307,52 +196,27 @@ def delete_file(remote_path: str) -> bool:
         logger.error(f"❌ Failed to delete file {remote_path}: {e}")
         raise IOError(f"Deletion failed: {e}")
 
-# ------------------------------------------------------------------------
-# Bulk Operations
-# ------------------------------------------------------------------------
+def file_exists(remote_path: str) -> bool:
+    _validate_path(remote_path)
+    return _fs.exists(_bucket_path(remote_path))
+
 def list_files(prefix: str = "", recursive: bool = True) -> List[str]:
-    """
-    List files in the bucket under a given prefix.
-    
-    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 = _bucket_path(prefix) if prefix else BUCKET_URI
-    
     try:
         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
+        prefix_len = len(BUCKET_URI) + 1
         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}")
+        logger.error(f"Failed to list files: {e}")
         return []
 
-def file_exists(remote_path: str) -> bool:
-    """Check if a file exists in the bucket."""
-    _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.
-    
-    Returns:
-        File size in bytes, or None if file doesn't exist
-    """
     _validate_path(remote_path)
-    
     try:
         info = _fs.info(_bucket_path(remote_path))
         return info.get("size")
@@ -363,14 +227,7 @@ def get_file_size(remote_path: str) -> Optional[int]:
         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 {
@@ -386,60 +243,25 @@ def get_file_info(remote_path: str) -> Optional[Dict[str, Any]]:
         logger.error(f"Failed to get file info: {e}")
         return None
 
-# ------------------------------------------------------------------------
-# Stream-based Operations
-# ------------------------------------------------------------------------
 def store_file_stream(remote_path: str, data_stream: BinaryIO, encrypt: bool = True) -> str:
-    """
-    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, 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)
 
-# ------------------------------------------------------------------------
-# Storage Statistics
-# ------------------------------------------------------------------------
 def get_storage_stats() -> Dict[str, Any]:
-    """
-    Get storage statistics.
-    
-    Returns:
-        Dict with storage information
-    """
     try:
         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 {
             "bucket": f"{OWNER}/{BUCKET_ID}",
             "file_count": file_count,
@@ -456,22 +278,9 @@ def get_storage_stats() -> Dict[str, Any]:
             "error": str(e)
         }
 
-# ------------------------------------------------------------------------
-# 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:
@@ -481,43 +290,22 @@ def create_backup(backup_prefix: str = "backups") -> str:
                 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."""
     if _fs:
         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'
+    'store_file', 'retrieve_file', 'delete_file', 'file_exists',
+    'list_files', 'get_file_size', 'get_file_info',
+    'store_file_stream', 'retrieve_file_stream',
+    'get_storage_stats', 'create_backup', 'close'
 ]