Spaces:
Sleeping
Sleeping
File size: 3,138 Bytes
c2e5995 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import os
import uuid
import logging
from pathlib import Path
from typing import Tuple, Optional
# Configure Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Detect environment (Hugging Face Spaces vs Local)
# HF Spaces with persistent storage usually mount at /data
IS_HF_SPACE = os.path.exists('/data')
if IS_HF_SPACE:
BASE_STORAGE_DIR = Path('/data/storage')
logger.info(f"Using PERSISTENT storage at {BASE_STORAGE_DIR}")
else:
BASE_STORAGE_DIR = Path(os.path.dirname(os.path.abspath(__file__))) / "storage"
logger.info(f"Using LOCAL storage at {BASE_STORAGE_DIR}")
def get_user_storage_path(username: str) -> Path:
"""Get secure storage path for user, creating it if needed."""
# Sanitize username to prevent directory traversal
safe_username = "".join([c for c in username if c.isalnum() or c in ('-', '_')])
user_path = BASE_STORAGE_DIR / safe_username
user_path.mkdir(parents=True, exist_ok=True)
return user_path
def save_image(username: str, file_bytes: bytes, filename_hint: str = "image.png") -> str:
"""
Save image to disk and return a unique image_id.
Returns: image_id (e.g. IMG_ABC123)
"""
# Generate ID
unique_suffix = uuid.uuid4().hex[:12].upper()
image_id = f"IMG_{unique_suffix}"
# Determine extension
ext = os.path.splitext(filename_hint)[1].lower()
if not ext:
ext = ".png" # Default
filename = f"{image_id}{ext}"
user_path = get_user_storage_path(username)
file_path = user_path / filename
try:
with open(file_path, "wb") as f:
f.write(file_bytes)
logger.info(f"Saved image {image_id} for user {username} at {file_path}")
return image_id
except Exception as e:
logger.error(f"Failed to save image: {e}")
raise IOError(f"Storage Error: {e}")
def load_image(username: str, image_id: str) -> Tuple[bytes, str]:
"""
Load image bytes from disk.
Returns: (file_bytes, file_path_str)
"""
# Security: Ensure ID format is valid
if not image_id.startswith("IMG_") or ".." in image_id or "/" in image_id:
raise ValueError("Invalid image_id format")
user_path = get_user_storage_path(username)
# We don't know the extension, so look for the file
# Or strict requirement: user must know?
# Better: Search for matching file
for file in user_path.glob(f"{image_id}.*"):
try:
with open(file, "rb") as f:
return f.read(), str(file)
except Exception as e:
logger.error(f"Error reading file {file}: {e}")
raise IOError("Read error")
raise FileNotFoundError(f"Image {image_id} not found for user {username}")
def get_image_absolute_path(username: str, image_id: str) -> Optional[str]:
"""Return absolute path if exists, else None."""
user_path = get_user_storage_path(username)
for file in user_path.glob(f"{image_id}.*"):
return str(file)
return None
|