Create storage_handler.py

storage_handler.py

@@ -0,0 +1,105 @@
+import os
+import base64
+import tempfile
+from typing import Optional, BinaryIO
+import logging
+
+from huggingface_hub import HfFileSystem, HfApi, create_bucket
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+logger = logging.getLogger("StorageHandler")
+
+# ------------------------------------------------------------------------
+# Configuration
+# ------------------------------------------------------------------------
+HF_TOKEN = os.environ["HF_TOKEN"]  # 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:
+    # Auto‑generate (for production, set env var permanently)
+    from cryptography.fernet import Fernet
+    FILE_ENCRYPTION_KEY_B64 = Fernet.generate_key().decode()
+FILE_ENCRYPTION_KEY = base64.urlsafe_b64decode(FILE_ENCRYPTION_KEY_B64.encode())
+assert len(FILE_ENCRYPTION_KEY) == 32, "AES‑256 requires a 32‑byte key"
+
+# Fetch the bucket owner (username) from the token
+_api = HfApi(token=HF_TOKEN)
+OWNER = _api.whoami()["name"]
+del _api
+
+# Bucket URI prefix
+BUCKET_URI_PREFIX = f"hf://buckets/{OWNER}/{BUCKET_NAME}"
+
+# ------------------------------------------------------------------------
+# Bucket lifecycle
+# ------------------------------------------------------------------------
+def ensure_bucket():
+    """Create the private bucket if it does not exist (idempotent)."""
+    try:
+        create_bucket(name=BUCKET_NAME, token=HF_TOKEN, private=True, exist_ok=True)
+        logger.info(f"Bucket '{BUCKET_NAME}' ready (owner: {OWNER}).")
+    except Exception as e:
+        logger.error(f"Failed to ensure bucket: {e}")
+        raise
+
+# Initialise on import
+ensure_bucket()
+
+# ------------------------------------------------------------------------
+# Encryption helpers
+# ------------------------------------------------------------------------
+def encrypt_file(data: bytes) -> bytes:
+    """AES‑256‑GCM: nonce (12 B) || ciphertext (incl. tag)"""
+    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:
+    nonce = encrypted_blob[:12]
+    ciphertext = encrypted_blob[12:]
+    aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
+    return aesgcm.decrypt(nonce, ciphertext, None)
+
+# ------------------------------------------------------------------------
+# Core file operations via HfFileSystem
+# ------------------------------------------------------------------------
+_fs = HfFileSystem(token=HF_TOKEN)
+
+def _bucket_path(remote_path: str) -> str:
+    """Constructs the full hf://buckets/... URI."""
+    return f"{BUCKET_URI_PREFIX}/{remote_path}"
+
+def store_file(remote_path: str, data: bytes) -> str:
+    """Encrypt and write data directly into the bucket. Returns the remote path."""
+    encrypted = encrypt_file(data)
+    full_uri = _bucket_path(remote_path)
+    with _fs.open(full_uri, "wb") as f:
+        f.write(encrypted)
+    return remote_path
+
+def retrieve_file(remote_path: str) -> bytes:
+    """Read from bucket and decrypt."""
+    full_uri = _bucket_path(remote_path)
+    with _fs.open(full_uri, "rb") as f:
+        encrypted = f.read()
+    return decrypt_file(encrypted)
+
+def delete_file(remote_path: str):
+    """Remove a file from the bucket."""
+    full_uri = _bucket_path(remote_path)
+    try:
+        _fs.rm(full_uri)
+    except FileNotFoundError:
+        pass
+
+# ------------------------------------------------------------------------
+# Helper for reassembled chunked uploads (used by app.py)
+# ------------------------------------------------------------------------
+def store_file_stream(remote_path: str, data_stream: BinaryIO) -> str:
+    """
+    Read the entire stream, encrypt, and write to bucket.
+    Useful for handling reassembled chunks without a single byte array.
+    """
+    data = data_stream.read()
+    return store_file(remote_path, data)