droplake/backend/routes/upload.py

import os
import secrets
import string
import struct
from datetime import datetime, timedelta
from pathlib import Path

from argon2 import PasswordHasher
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
from sqlalchemy.orm import Session

from config import settings
from database import Share, SecurityEvent, get_db

router = APIRouter()

_ph = PasswordHasher(time_cost=2, memory_cost=65536, parallelism=2, hash_len=32, salt_len=16)

# ── Constants ──────────────────────────────────────────────────────────────

CHUNK_SIZE   = 64 * 1024 * 1024          # 64 MB plaintext per chunk
BASE62       = string.ascii_letters + string.digits
VALID_TTL    = {1, 6, 24, 48, 72}        # hours; 72h max — ephemeral sharing, not storage

# ── Helpers ────────────────────────────────────────────────────────────────

def _share_id() -> str:
    """10-character base62 ID — ~59 bits of entropy, brute-force safe with rate limiting."""
    return ''.join(secrets.choice(BASE62) for _ in range(10))

def _sanitize_filename(name: str) -> str:
    safe = Path(name).name
    safe = ''.join(c for c in safe if c.isprintable() and c not in r'\/:*?"<>|')
    return safe[:255] or 'unnamed'

async def _encrypt_to_disk(upload: UploadFile, out_path: Path, key: bytes) -> int:
    """
    Stream-encrypt upload into chunked AES-256-GCM format.

    File layout:
      [4B uint32 BE — chunk count]
      Per chunk:
        [4B uint32 BE — len(nonce + ciphertext)]
        [12B nonce]
        [N B ciphertext + 16B GCM tag]

    Returns plaintext byte count.
    """
    aesgcm = AESGCM(key)
    chunk_count = 0
    plaintext_bytes = 0
    buf = bytearray()

    try:
        with open(out_path, 'wb') as f:
            f.write(b'\x00\x00\x00\x00')           # placeholder for chunk count

            while True:
                data = await upload.read(1024 * 1024)   # 1 MB reads
                if not data:
                    break
                plaintext_bytes += len(data)
                if plaintext_bytes > settings.max_upload_bytes:
                    raise HTTPException(413, 'Upload exceeds size limit')
                buf.extend(data)

                while len(buf) >= CHUNK_SIZE:
                    chunk = bytes(buf[:CHUNK_SIZE])
                    del buf[:CHUNK_SIZE]
                    nonce = os.urandom(12)
                    ct = aesgcm.encrypt(nonce, chunk, None)
                    payload = nonce + ct
                    f.write(struct.pack('>I', len(payload)))
                    f.write(payload)
                    chunk_count += 1

            # Final partial chunk
            if buf:
                nonce = os.urandom(12)
                ct = aesgcm.encrypt(nonce, bytes(buf), None)
                payload = nonce + ct
                f.write(struct.pack('>I', len(payload)))
                f.write(payload)
                chunk_count += 1

            # Write real chunk count
            f.seek(0)
            f.write(struct.pack('>I', chunk_count))

    except HTTPException:
        out_path.unlink(missing_ok=True)
        raise
    except Exception:
        out_path.unlink(missing_ok=True)
        raise HTTPException(500, 'Encryption failed during upload')

    return plaintext_bytes

# ── Route ──────────────────────────────────────────────────────────────────

@router.post('/upload')
async def upload_share(
    request: Request,
    file: UploadFile = File(...),
    share_type: str = Form(...),
    original_filename: str = Form(None),
    mimetype: str = Form(None),
    ttl_hours: int = Form(24),
    burn_after_read: bool = Form(False),
    password: str = Form(None),
    db: Session = Depends(get_db),
):
    if share_type not in ('file', 'text'):
        raise HTTPException(400, 'Invalid share_type')
    if ttl_hours not in VALID_TTL:
        raise HTTPException(400, 'Invalid TTL value')
    if password and len(password) > 128:
        raise HTTPException(400, 'Password too long')

    if original_filename:
        original_filename = _sanitize_filename(original_filename)
    if mimetype:
        mimetype = mimetype[:128]

    # Generate key and storage path
    key        = os.urandom(32)                  # AES-256 key
    share_id   = _share_id()
    blob_name  = secrets.token_hex(16)           # random storage filename
    blob_path  = Path(settings.storage_path) / blob_name
    Path(settings.storage_path).mkdir(parents=True, exist_ok=True)

    plaintext_bytes = await _encrypt_to_disk(file, blob_path, key)

    password_hash = _ph.hash(password) if password else None

    now = datetime.utcnow()
    share = Share(
        id               = share_id,
        share_type       = share_type,
        original_filename= original_filename,
        mimetype         = mimetype,
        filesize         = plaintext_bytes,
        storage_path     = blob_name,
        encryption_key   = key.hex(),
        password_hash    = password_hash,
        created_at       = now,
        expires_at       = now + timedelta(hours=ttl_hours),
        burn_after_read  = burn_after_read,
    )
    db.add(share)
    db.add(SecurityEvent(event_type='share_created', share_id=share_id))
    db.commit()

    return {
        'share_id'  : share_id,
        'expires_at': share.expires_at.isoformat() + 'Z',
        'has_password': password_hash is not None,
    }