Spaces:
Running
Running
| 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' | |
| ] |