Create app/utils.py

app/utils.py

@@ -0,0 +1,125 @@
+# app/utils.py
+import os
+import json
+import time
+import asyncio
+from pathlib import Path
+from typing import Dict, Set
+
+from fastapi import WebSocket
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.backends import default_backend
+
+# --- Constants ---
+UPLOAD_ROOT = Path("/tmp/triflix_uploader")
+# In a real app, load this securely from env/secrets management
+try:
+    MASTER_KEY = bytes.fromhex(os.environ["TRIFLIX_MASTER_KEY"])
+    if len(MASTER_KEY) != 32:
+        raise ValueError("TRIFLIX_MASTER_KEY must be a 32-byte (64-char) hex string.")
+except (KeyError, ValueError) as e:
+    print(f"WARNING: Invalid or missing TRIFLIX_MASTER_KEY. Using a temporary key for demo purposes. {e}")
+    MASTER_KEY = os.urandom(32)
+
+# --- In-memory State (locks and WebSocket connections) ---
+# These are reset if the server restarts. The on-disk state persists.
+session_locks: Dict[str, asyncio.Lock] = {}
+ws_connections: Dict[str, Set[WebSocket]] = {}
+
+# --- Session Management ---
+def get_session_dir(session_id: str) -> Path:
+    """Returns the directory for a given session."""
+    return UPLOAD_ROOT / session_id
+
+def load_meta(session_id: str) -> dict:
+    """Loads metadata for a session from its meta.json file."""
+    meta_path = get_session_dir(session_id) / "meta.json"
+    if not meta_path.exists():
+        raise FileNotFoundError("Session metadata not found.")
+    return json.loads(meta_path.read_text())
+
+def save_meta(session_id: str, meta: dict):
+    """Saves metadata for a session to its meta.json file."""
+    meta_path = get_session_dir(session_id) / "meta.json"
+    meta_path.parent.mkdir(parents=True, exist_ok=True)
+    meta_path.write_text(json.dumps(meta, indent=2))
+
+async def broadcast_progress(session_id: str):
+    """Sends progress updates to all connected WebSocket clients for a session."""
+    if session_id not in ws_connections:
+        return
+        
+    try:
+        meta = load_meta(session_id)
+        payload = {
+            "type": "progress",
+            "uploaded_bytes": meta.get("uploaded_bytes", 0),
+            "total_bytes": meta.get("total_bytes"),
+            "status": meta.get("status", "uploading"),
+        }
+        # Create a list of tasks to send messages concurrently
+        tasks = [ws.send_json(payload) for ws in ws_connections[session_id]]
+        await asyncio.gather(*tasks, return_exceptions=True) # Use gather to handle potential disconnects
+    except FileNotFoundError:
+        # If meta is gone, maybe the session was cleaned up. Nothing to broadcast.
+        pass
+
+# --- Rate Limiting ---
+rate_limit_store: Dict[str, list] = {}
+
+def enforce_rate_limit(ip: str, limit: int = 10, per_seconds: int = 60) -> bool:
+    """Simple in-memory rate limiter. Returns True if allowed, False if denied."""
+    now = time.time()
+    timestamps = rate_limit_store.setdefault(ip, [])
+    # Remove timestamps older than the window
+    timestamps[:] = [t for t in timestamps if t > now - per_seconds]
+    if len(timestamps) >= limit:
+        return False
+    timestamps.append(now)
+    return True
+
+# --- Streaming Encryption/Decryption ---
+CHUNK_SIZE = 64 * 1024  # 64KB
+
+def encrypt_file(source_path: Path, dest_path: Path):
+    """Encrypts a file using AES-GCM streaming."""
+    nonce = os.urandom(12)
+    cipher = Cipher(algorithms.AES(MASTER_KEY), modes.GCM(nonce), backend=default_backend()).encryptor()
+    
+    with open(source_path, "rb") as fin, open(dest_path, "wb") as fout:
+        fout.write(nonce)
+        while True:
+            chunk = fin.read(CHUNK_SIZE)
+            if not chunk:
+                break
+            encrypted_chunk = cipher.update(chunk)
+            fout.write(encrypted_chunk)
+        
+        fout.write(cipher.finalize())
+        fout.write(cipher.tag)
+
+async def decrypt_stream_generator(encrypted_path: Path):
+    """A generator that yields decrypted chunks of a file."""
+    with open(encrypted_path, "rb") as f:
+        nonce = f.read(12)
+        # Read the rest of the file to get ciphertext and tag
+        f.seek(0, os.SEEK_END)
+        end_pos = f.tell()
+        f.seek(12) # Go back to after the nonce
+        
+        ciphertext_with_tag = f.read()
+        ciphertext = ciphertext_with_tag[:-16]
+        tag = ciphertext_with_tag[-16:]
+        
+        decryptor = Cipher(algorithms.AES(MASTER_KEY), modes.GCM(nonce, tag), backend=default_backend()).decryptor()
+        
+        # We can process the ciphertext in chunks now
+        offset = 0
+        while offset < len(ciphertext):
+            chunk = ciphertext[offset:offset + CHUNK_SIZE]
+            decrypted_chunk = decryptor.update(chunk)
+            yield decrypted_chunk
+            offset += CHUNK_SIZE
+        
+        # Finalize to check tag and get any remaining data
+        yield decryptor.finalize()