Spaces:
Sleeping
Sleeping
| # storage_handler.py | |
| import os | |
| import base64 | |
| import logging | |
| from typing import Optional, BinaryIO, Dict, Any | |
| from huggingface_hub import HfFileSystem, HfApi | |
| from cryptography.hazmat.primitives.ciphers.aead import AESGCM | |
| logger = logging.getLogger("StorageHandler") | |
| # ------------------------------------------------------------------------ | |
| # Configuration – all from environment variables | |
| # ------------------------------------------------------------------------ | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| if not HF_TOKEN: | |
| raise ValueError("HF_TOKEN environment variable is required") | |
| BUCKET_ID = 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" | |
| # Get the bucket owner (your 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'") | |
| # Build the proper bucket URI prefix (Xet-native) | |
| BUCKET_URI_PREFIX = f"hf://buckets/{OWNER}/{BUCKET_ID}" | |
| # ------------------------------------------------------------------------ | |
| # HfFileSystem for low-level bucket operations (Xet-native, NO Git/LFS) | |
| # ------------------------------------------------------------------------ | |
| _fs = HfFileSystem(token=HF_TOKEN) | |
| # ------------------------------------------------------------------------ | |
| # Bucket initialization (using the modern Hugging Face Hub API) | |
| # ------------------------------------------------------------------------ | |
| 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) | |
| 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: | |
| # 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 | |
| else: | |
| try: | |
| 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}" | |
| except Exception as e2: | |
| if "already exists" in str(e2).lower(): | |
| logger.info(f"Bucket '{BUCKET_ID}' already exists.") | |
| else: | |
| logger.error(f"Failed to create bucket: {e2}") | |
| raise | |
| # Initialize bucket | |
| ensure_bucket() | |
| # ------------------------------------------------------------------------ | |
| # Encryption / Decryption | |
| # ------------------------------------------------------------------------ | |
| def encrypt_file(data: bytes) -> bytes: | |
| """Encrypt binary data using AES‑256‑GCM. Output: nonce (12 B) || ciphertext.""" | |
| aesgcm = AESGCM(FILE_ENCRYPTION_KEY) | |
| nonce = os.urandom(12) | |
| ciphertext = aesgcm.encrypt(nonce, data, None) | |
| return nonce + ciphertext | |
| def decrypt_file(encrypted_blob: bytes) -> bytes: | |
| """Decrypt data produced by encrypt_file().""" | |
| nonce = encrypted_blob[:12] | |
| ciphertext = encrypted_blob[12:] | |
| aesgcm = AESGCM(FILE_ENCRYPTION_KEY) | |
| return aesgcm.decrypt(nonce, ciphertext, None) | |
| # ------------------------------------------------------------------------ | |
| # Core Bucket Operations (Xet-native via HfFileSystem) | |
| # ------------------------------------------------------------------------ | |
| def _bucket_path(remote_path: str) -> str: | |
| """Construct the full hf://buckets/... URI.""" | |
| return f"{BUCKET_URI_PREFIX}/{remote_path}" | |
| def store_file(remote_path: str, data: bytes) -> str: | |
| """ | |
| Encrypt data and write it to the bucket using HfFileSystem. | |
| Returns: the remote path (for DB storage). | |
| """ | |
| encrypted_data = encrypt_file(data) | |
| full_uri = _bucket_path(remote_path) | |
| with _fs.open(full_uri, "wb") as f: | |
| f.write(encrypted_data) | |
| logger.debug(f"Stored file: {full_uri}") | |
| return remote_path | |
| def retrieve_file(remote_path: str) -> bytes: | |
| """ | |
| Read from bucket using HfFileSystem and decrypt. | |
| """ | |
| full_uri = _bucket_path(remote_path) | |
| with _fs.open(full_uri, "rb") as f: | |
| encrypted_data = f.read() | |
| return decrypt_file(encrypted_data) | |
| def delete_file(remote_path: str): | |
| """ | |
| Delete a file from the bucket using HfFileSystem. | |
| """ | |
| 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}") | |
| except Exception as e: | |
| logger.error(f"Failed to delete file {full_uri}: {e}") | |
| def list_files(prefix: str = "") -> list: | |
| """ | |
| List files in the bucket under a given prefix. | |
| Returns list of full URIs. | |
| """ | |
| search_path = f"{BUCKET_URI_PREFIX}/{prefix}" if prefix else BUCKET_URI_PREFIX | |
| try: | |
| return _fs.ls(search_path) | |
| except FileNotFoundError: | |
| 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) | |
| def get_file_size(remote_path: str) -> Optional[int]: | |
| """Get the size of a file in bytes (compressed).""" | |
| full_uri = _bucket_path(remote_path) | |
| try: | |
| info = _fs.info(full_uri) | |
| return info.get("size") | |
| except Exception: | |
| return None | |
| # ------------------------------------------------------------------------ | |
| # Stream-based store (for reassembled chunks) | |
| # ------------------------------------------------------------------------ | |
| def store_file_stream(remote_path: str, data_stream: BinaryIO) -> str: | |
| """ | |
| Read the entire stream, encrypt, and write to bucket. | |
| """ | |
| data = data_stream.read() | |
| return store_file(remote_path, data) | |
| # ------------------------------------------------------------------------ | |
| # Bucket info (for monitoring / health checks) | |
| # ------------------------------------------------------------------------ | |
| def get_bucket_info() -> Dict[str, Any]: | |
| """Return information about the bucket.""" | |
| api = HfApi(token=HF_TOKEN) | |
| try: | |
| # Try the modern API first | |
| info = api.bucket_info(bucket_id=f"{OWNER}/{BUCKET_ID}") | |
| return { | |
| "id": f"{OWNER}/{BUCKET_ID}", | |
| "exists": True, | |
| "private": info.get("private", True), | |
| "total_files": len(list_files()), | |
| } | |
| except Exception: | |
| return { | |
| "id": f"{OWNER}/{BUCKET_ID}", | |
| "exists": False, | |
| "note": "Cannot fetch bucket info" | |
| } | |
| # ------------------------------------------------------------------------ | |
| # Clean shutdown | |
| # ------------------------------------------------------------------------ | |
| def close(): | |
| """Clean up resources (HfFileSystem).""" | |
| if _fs: | |
| _fs.close() |