simple-nanoshare/routes/side/upload.py

from my_modules.decoratory.header import login_required
from my_modules.expiry import parse_expires, ensure_utc
from my_modules.file_meta import iso_stamp_filename, format_size

from quart import Blueprint, request, jsonify, current_app
import asyncio, hashlib

upload_bp = Blueprint('upload_bp', __name__)

# --- Helpers -----------------------------------------------------

async def read_all(uploaded) -> bytes:
  """Read all bytes from an uploaded file, handling sync or async .read()."""
  reader = getattr(uploaded, 'read', None)
  if reader is None:
    return b''
  if asyncio.iscoroutinefunction(reader):
    return await reader()

  data = reader()
  if asyncio.iscoroutine(data):
    return await data
  return data


async def fingerprint_stream(stream, chunk_size:int=1024 * 1024) -> tuple[str|None, int|None]:
  if not hasattr(stream, 'seek') or not hasattr(stream, 'tell'):
    return None, None

  try:
    stream.seek(0)
  except Exception:
    return None, None

  digest = hashlib.sha256()
  size_bytes = 0

  while True:
    chunk = await asyncio.to_thread(stream.read, chunk_size)
    if not chunk:
      break
    size_bytes += len(chunk)
    digest.update(chunk)

  try:
    stream.seek(0)
  except Exception:
    return None, None

  return digest.hexdigest(), size_bytes


def fingerprint_bytes(data: bytes) -> str:
  return hashlib.sha256(data).hexdigest()


# --- Routes ------------------------------------------------------

@upload_bp.post('/api/upload')
@login_required
async def api_upload(user):
  """
  POST /upload/api/upload
  Accepts:
    - multipart form with 'file' or 'text'
    - 'expires' can be '1h', '7d', or ISO timestamp
    - 'note' optional
  """
  form = await request.form
  files = await request.files
  note = form.get('note', '')
  expires_raw = form.get('expires', '')
  text = form.get('text', '')
  orphan_registry = getattr(current_app, 'orphan_storage_registry', None)

  uploaded = files.get('file')
  expires_at_dt = ensure_utc(parse_expires(expires_raw))

  if not uploaded and not text.strip():
    return jsonify({'ok': False, 'error': 'No content provided'}), 400

  content_type = None

  # --- binary upload path ---
  if uploaded:
    fname = uploaded.filename or ''
    ctype = uploaded.mimetype or 'application/octet-stream'
    content_type = ctype
    storage_id = None
    size_bytes = 0
    fingerprint = None
    reused_orphan_storage_id = False

    # generate filename if missing/placeholder
    if not fname or fname.lower() in {'blob', 'file'}:
      ext = {
        'image/png': 'png',
        'image/jpeg': 'jpg',
        'image/gif': 'gif',
        'image/webp': 'webp',
        'application/pdf': 'pdf',
        'text/plain': 'txt',
      }.get(ctype, 'bin')
      fname = iso_stamp_filename('pasted', ext)

    stream = getattr(uploaded, 'stream', None)

    if stream is not None:
      fingerprint, detected_size = await fingerprint_stream(stream)
      size_bytes = detected_size or 0
      storage_id = (
        await orphan_registry.pop_recent(user['sub'], fingerprint)
        if orphan_registry
        else None
      )
      if storage_id:
        reused_orphan_storage_id = True
      else:
        storage_id, sent_size = await current_app.convex.send_stream_to_storage(stream=stream,content_type=content_type)
        size_bytes = sent_size
    else:
      data = await read_all(uploaded)
      fingerprint = fingerprint_bytes(data)
      storage_id = (
        await orphan_registry.pop_recent(user['sub'], fingerprint)
        if orphan_registry
        else None
      )
      if storage_id:
        reused_orphan_storage_id = True
      else:
        storage_id = await current_app.convex.send_to_storage(data=data, content_type=content_type)
      size_bytes = len(data)

    file_size_pretty = format_size(size_bytes)

    try:
      await current_app.convex.add_file(
        file_name=fname,
        file_size=file_size_pretty,
        note=note,
        content_type=content_type,
        expires_at=expires_at_dt,
        storage_id=storage_id,
        user_id=user['sub'],
      )
    except Exception:
      if storage_id and not reused_orphan_storage_id and orphan_registry:
        await orphan_registry.remember(user['sub'], fingerprint, storage_id)
      raise

  # --- text upload path ---
  elif text.strip():
    data = text.encode('utf-8')
    fname = iso_stamp_filename('pasted', 'txt')
    fingerprint = fingerprint_bytes(data)
    storage_id = (
      await orphan_registry.pop_recent(user['sub'], fingerprint)
      if orphan_registry
      else None
    )
    reused_orphan_storage_id = bool(storage_id)
    if not storage_id:
      storage_id = await current_app.convex.send_to_storage(
        data=data, content_type='text/plain'
      )

    size_bytes = len(data)
    file_size_pretty = format_size(size_bytes)

    try:
      await current_app.convex.add_file(
        file_name=fname,
        file_size=file_size_pretty,
        note=note,
        content_type='text/plain',
        expires_at=expires_at_dt,
        storage_id=storage_id,
        user_id=user['sub'],
      )
    except Exception:
      if not reused_orphan_storage_id and orphan_registry:
        await orphan_registry.remember(user['sub'], fingerprint, storage_id)
      raise

  return jsonify({'ok': True})