storage_handler.py

import os
import base64
import hashlib
import uuid
import logging
from typing import Optional, BinaryIO, Dict, Any, List
import time

from huggingface_hub import HfFileSystem, HfApi
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

logger = logging.getLogger("InfinityChat.Storage")

# ------------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------------
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("❌ HF_TOKEN environment variable is required")

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!")
    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 bucket owner (HF username)
# ------------------------------------------------------------------------
_api = HfApi(token=HF_TOKEN)
try:
    owner_info = _api.whoami()
    OWNER = owner_info["name"]
except Exception as e:
    logger.warning(f"Could not determine HF username: {e}")
    OWNER = "infinitychat"

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 instance
# ------------------------------------------------------------------------
_fs = HfFileSystem(token=HF_TOKEN)

# ------------------------------------------------------------------------
# Bucket Initialization
# ------------------------------------------------------------------------
def ensure_bucket():
    """Create the private bucket if it doesn't exist."""
    global OWNER, BUCKET_URI, BUCKET_ID

    try:
        _api.create_bucket(
            bucket_id=f"{OWNER}/{BUCKET_ID}",
            private=True,
            exist_ok=True
        )
        logger.info(f"✅ Bucket '{OWNER}/{BUCKET_ID}' is ready (private)")
    except Exception as e:
        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:
                _api.create_bucket(
                    bucket_id=BUCKET_ID,
                    private=True,
                    exist_ok=True
                )
                OWNER = ""
                BUCKET_URI = f"hf://buckets/{BUCKET_ID}"
                logger.info(f"✅ Bucket '{BUCKET_ID}' is ready (alternative format)")
            except Exception as e2:
                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}")
                    logger.info("The app will still work but file uploads will fail")

ensure_bucket()

# ------------------------------------------------------------------------
# Encryption / Decryption
# ------------------------------------------------------------------------
def encrypt_bytes(data: bytes, aad: Optional[bytes] = None) -> bytes:
    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:
    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:
    return len(encrypted_blob) >= 29

# ------------------------------------------------------------------------
# Path Utilities
# ------------------------------------------------------------------------
def _build_path(*parts: str) -> str:
    return "/".join(p.strip("/") for p in parts if p)

def _bucket_path(remote_path: str) -> str:
    return f"{BUCKET_URI}/{remote_path.lstrip('/')}"

def _validate_path(remote_path: str) -> None:
    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:
    _validate_path(remote_path)
    try:
        if encrypt:
            encrypted_data = encrypt_bytes(data, remote_path.encode('utf-8'))
        else:
            encrypted_data = data
        full_uri = _bucket_path(remote_path)
        with _fs.open(full_uri, "wb") as f:
            f.write(encrypted_data)
        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:
    _validate_path(remote_path)
    full_uri = _bucket_path(remote_path)
    try:
        if not _fs.exists(full_uri):
            raise FileNotFoundError(f"File not found: {remote_path}")
        with _fs.open(full_uri, "rb") as f:
            encrypted_data = f.read()
        if not encrypted_data:
            raise IOError(f"Empty file: {remote_path}")
        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) -> bool:
    _validate_path(remote_path)
    full_uri = _bucket_path(remote_path)
    try:
        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 {remote_path}: {e}")
        raise IOError(f"Deletion failed: {e}")

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]:
    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)
        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}")
        return []

def get_file_size(remote_path: str) -> Optional[int]:
    _validate_path(remote_path)
    try:
        info = _fs.info(_bucket_path(remote_path))
        return info.get("size")
    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]]:
    _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

def store_file_stream(remote_path: str, data_stream: BinaryIO, encrypt: bool = True) -> str:
    data = data_stream.read()
    return store_file(remote_path, data, encrypt=encrypt)

def retrieve_file_stream(remote_path: str, decrypt: bool = True) -> BinaryIO:
    data = retrieve_file(remote_path, decrypt=decrypt)
    import io
    return io.BytesIO(data)

def get_storage_stats() -> Dict[str, Any]:
    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,
            "total_size": total_size,
            "total_size_mb": round(total_size / (1024 * 1024), 2)
        }
    except Exception as e:
        logger.error(f"Failed to get storage stats: {e}")
        return {
            "bucket": f"{OWNER}/{BUCKET_ID}",
            "file_count": 0,
            "total_size": 0,
            "total_size_mb": 0,
            "error": str(e)
        }

def create_backup(backup_prefix: str = "backups") -> str:
    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

def close():
    if _fs:
        try:
            _fs.close()
        except:
            pass

__all__ = [
    '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'
]