RC1
This commit is contained in:
oib
266
upload.py
266
upload.py
@ -6,11 +6,13 @@ from slowapi.util import get_remote_address
|
||||
from slowapi.errors import RateLimitExceeded
|
||||
from pathlib import Path
|
||||
import json
|
||||
import requests
|
||||
from datetime import datetime
|
||||
from convert_to_opus import convert_to_opus
|
||||
from models import UploadLog, UserQuota, User
|
||||
from sqlalchemy import select
|
||||
from models import UploadLog, UserQuota, User, PublicStream
|
||||
from sqlalchemy import select, or_
|
||||
from database import get_db
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
limiter = Limiter(key_func=get_remote_address)
|
||||
router = APIRouter()
|
||||
@ -23,55 +25,63 @@ DATA_ROOT = Path("./data")
|
||||
@router.post("/upload")
|
||||
async def upload(request: Request, db = Depends(get_db), uid: str = Form(...), file: UploadFile = Form(...)):
|
||||
from log import log_violation
|
||||
import time
|
||||
|
||||
# Generate a unique request ID for this upload
|
||||
request_id = str(int(time.time()))
|
||||
log_violation("UPLOAD", request.client.host, uid, f"[{request_id}] Starting upload of {file.filename}")
|
||||
|
||||
try:
|
||||
user_dir = DATA_ROOT / uid
|
||||
user_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
raw_path = user_dir / ("raw." + file.filename.split(".")[-1])
|
||||
import uuid
|
||||
|
||||
unique_name = str(uuid.uuid4()) + ".opus"
|
||||
|
||||
# Save temp upload FIRST
|
||||
with open(raw_path, "wb") as f:
|
||||
f.write(await file.read())
|
||||
|
||||
# Block music/singing via Ollama prompt
|
||||
import requests
|
||||
try:
|
||||
with open(raw_path, "rb") as f:
|
||||
audio = f.read()
|
||||
res = requests.post("http://localhost:11434/api/generate", json={
|
||||
"model": "whisper",
|
||||
"prompt": "Does this audio contain music or singing? Answer yes or no only.",
|
||||
"audio": audio
|
||||
}, timeout=10)
|
||||
resp = res.json().get("response", "").lower()
|
||||
if "yes" in resp:
|
||||
raw_path.unlink(missing_ok=True)
|
||||
raise HTTPException(status_code=403, detail="Upload rejected: music or singing detected")
|
||||
except Exception as ollama_err:
|
||||
# fallback: allow, log if needed
|
||||
pass
|
||||
processed_path = user_dir / unique_name
|
||||
|
||||
# Block unconfirmed users (use ORM)
|
||||
# First, verify the user exists and is confirmed
|
||||
user = db.exec(select(User).where((User.username == uid) | (User.email == uid))).first()
|
||||
# If result is a Row or tuple, extract the User object
|
||||
if user is not None and not isinstance(user, User) and hasattr(user, "__getitem__"):
|
||||
user = user[0]
|
||||
from log import log_violation
|
||||
log_violation("UPLOAD", request.client.host, uid, f"DEBUG: Incoming uid={uid}, user found={user}, confirmed={getattr(user, 'confirmed', None)}")
|
||||
log_violation("UPLOAD", request.client.host, uid, f"DEBUG: After unpack, user={user}, type={type(user)}, confirmed={getattr(user, 'confirmed', None)}")
|
||||
|
||||
log_violation("UPLOAD", request.client.host, uid, f"[{request_id}] User check - found: {user is not None}, confirmed: {getattr(user, 'confirmed', False) if user else 'N/A'}")
|
||||
|
||||
if not user or not hasattr(user, "confirmed") or not user.confirmed:
|
||||
raw_path.unlink(missing_ok=True)
|
||||
raise HTTPException(status_code=403, detail="Account not confirmed")
|
||||
|
||||
# DB-based quota check
|
||||
quota = db.get(UserQuota, uid)
|
||||
if quota and quota.storage_bytes >= 100 * 1024 * 1024:
|
||||
raw_path.unlink(missing_ok=True)
|
||||
# Check quota before doing any file operations
|
||||
quota = db.get(UserQuota, uid) or UserQuota(uid=uid, storage_bytes=0)
|
||||
if quota.storage_bytes >= 100 * 1024 * 1024:
|
||||
raise HTTPException(status_code=400, detail="Quota exceeded")
|
||||
|
||||
# Create user directory if it doesn't exist
|
||||
user_dir = DATA_ROOT / uid
|
||||
user_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Generate a unique filename for the processed file first
|
||||
import uuid
|
||||
unique_name = f"{uuid.uuid4()}.opus"
|
||||
raw_ext = file.filename.split(".")[-1].lower()
|
||||
raw_path = user_dir / ("raw." + raw_ext)
|
||||
processed_path = user_dir / unique_name
|
||||
|
||||
# Clean up any existing raw files first (except the one we're about to create)
|
||||
for old_file in user_dir.glob('raw.*'):
|
||||
try:
|
||||
if old_file != raw_path: # Don't delete the file we're about to create
|
||||
old_file.unlink(missing_ok=True)
|
||||
log_violation("UPLOAD", request.client.host, uid, f"[{request_id}] Cleaned up old file: {old_file}")
|
||||
except Exception as e:
|
||||
log_violation("UPLOAD_ERROR", request.client.host, uid, f"[{request_id}] Failed to clean up {old_file}: {e}")
|
||||
|
||||
# Save the uploaded file temporarily
|
||||
log_violation("UPLOAD", request.client.host, uid, f"[{request_id}] Saving temporary file to {raw_path}")
|
||||
try:
|
||||
with open(raw_path, "wb") as f:
|
||||
content = await file.read()
|
||||
if not content:
|
||||
raise ValueError("Uploaded file is empty")
|
||||
f.write(content)
|
||||
log_violation("UPLOAD", request.client.host, uid, f"[{request_id}] Successfully wrote {len(content)} bytes to {raw_path}")
|
||||
except Exception as e:
|
||||
log_violation("UPLOAD_ERROR", request.client.host, uid, f"[{request_id}] Failed to save {raw_path}: {e}")
|
||||
raise HTTPException(status_code=500, detail=f"Failed to save uploaded file: {e}")
|
||||
|
||||
# Ollama music/singing check is disabled for this release
|
||||
log_violation("UPLOAD", request.client.host, uid, f"[{request_id}] Ollama music/singing check is disabled")
|
||||
|
||||
try:
|
||||
convert_to_opus(str(raw_path), str(processed_path))
|
||||
@ -82,44 +92,96 @@ async def upload(request: Request, db = Depends(get_db), uid: str = Form(...), f
|
||||
original_size = raw_path.stat().st_size
|
||||
raw_path.unlink(missing_ok=True) # cleanup
|
||||
|
||||
# First, verify the file was created and has content
|
||||
if not processed_path.exists() or processed_path.stat().st_size == 0:
|
||||
raise HTTPException(status_code=500, detail="Failed to process audio file")
|
||||
|
||||
# Concatenate all .opus files in random order to stream.opus for public playback
|
||||
# This is now done after the file is in its final location with log ID
|
||||
from concat_opus import concat_opus_files
|
||||
try:
|
||||
concat_opus_files(user_dir, user_dir / "stream.opus")
|
||||
except Exception as e:
|
||||
# fallback: just use the latest processed file if concat fails
|
||||
import shutil
|
||||
stream_path = user_dir / "stream.opus"
|
||||
shutil.copy2(processed_path, stream_path)
|
||||
|
||||
# Create a log entry with the original filename
|
||||
log = UploadLog(
|
||||
uid=uid,
|
||||
ip=request.client.host,
|
||||
filename=file.filename, # Store original filename
|
||||
processed_filename=unique_name, # Store the processed filename
|
||||
size_bytes=original_size
|
||||
)
|
||||
db.add(log)
|
||||
db.commit()
|
||||
db.refresh(log)
|
||||
|
||||
# Rename the processed file to include the log ID for better tracking
|
||||
processed_with_id = user_dir / f"{log.id}_{unique_name}"
|
||||
processed_path.rename(processed_with_id)
|
||||
processed_path = processed_with_id
|
||||
|
||||
# Store updated quota
|
||||
def update_stream_opus():
|
||||
try:
|
||||
concat_opus_files(user_dir, user_dir / "stream.opus")
|
||||
except Exception as e:
|
||||
# fallback: just use the latest processed file if concat fails
|
||||
import shutil
|
||||
stream_path = user_dir / "stream.opus"
|
||||
shutil.copy2(processed_path, stream_path)
|
||||
log_violation("STREAM_UPDATE", request.client.host, uid,
|
||||
f"[fallback] Updated stream.opus with {processed_path}")
|
||||
|
||||
# We'll call this after the file is in its final location
|
||||
|
||||
# Get the final file size
|
||||
size = processed_path.stat().st_size
|
||||
quota = db.get(UserQuota, uid)
|
||||
if not quota:
|
||||
quota = UserQuota(uid=uid)
|
||||
db.add(quota)
|
||||
quota.storage_bytes += size
|
||||
db.commit()
|
||||
|
||||
# Update public streams list
|
||||
update_public_streams(uid, quota.storage_bytes)
|
||||
# Start a transaction
|
||||
try:
|
||||
# Create a log entry with the original filename
|
||||
log = UploadLog(
|
||||
uid=uid,
|
||||
ip=request.client.host,
|
||||
filename=file.filename, # Store original filename
|
||||
processed_filename=unique_name, # Store the processed filename
|
||||
size_bytes=size
|
||||
)
|
||||
db.add(log)
|
||||
db.flush() # Get the log ID without committing
|
||||
|
||||
# Rename the processed file to include the log ID for better tracking
|
||||
processed_with_id = user_dir / f"{log.id}_{unique_name}"
|
||||
if processed_path.exists():
|
||||
# First check if there's already a file with the same UUID but different prefix
|
||||
for existing_file in user_dir.glob(f"*_{unique_name}"):
|
||||
if existing_file != processed_path:
|
||||
log_violation("CLEANUP", request.client.host, uid,
|
||||
f"[UPLOAD] Removing duplicate file: {existing_file}")
|
||||
existing_file.unlink(missing_ok=True)
|
||||
|
||||
# Now do the rename
|
||||
if processed_path != processed_with_id:
|
||||
if processed_with_id.exists():
|
||||
processed_with_id.unlink(missing_ok=True)
|
||||
processed_path.rename(processed_with_id)
|
||||
processed_path = processed_with_id
|
||||
|
||||
# Only clean up raw.* files, not previously uploaded opus files
|
||||
for old_temp_file in user_dir.glob('raw.*'):
|
||||
try:
|
||||
old_temp_file.unlink(missing_ok=True)
|
||||
log_violation("CLEANUP", request.client.host, uid, f"[{request_id}] Cleaned up temp file: {old_temp_file}")
|
||||
except Exception as e:
|
||||
log_violation("CLEANUP_ERROR", request.client.host, uid, f"[{request_id}] Failed to clean up {old_temp_file}: {e}")
|
||||
|
||||
# Get or create quota
|
||||
quota = db.query(UserQuota).filter(UserQuota.uid == uid).first()
|
||||
if not quota:
|
||||
quota = UserQuota(uid=uid, storage_bytes=0)
|
||||
db.add(quota)
|
||||
|
||||
# Update quota with the new file size
|
||||
quota.storage_bytes = sum(
|
||||
f.stat().st_size
|
||||
for f in user_dir.glob('*.opus')
|
||||
if f.name != 'stream.opus' and f != processed_path
|
||||
) + size
|
||||
|
||||
# Update public streams
|
||||
update_public_streams(uid, quota.storage_bytes, db)
|
||||
|
||||
# Commit the transaction
|
||||
db.commit()
|
||||
|
||||
# Now that the transaction is committed and files are in their final location,
|
||||
# update the stream.opus file to include all files
|
||||
update_stream_opus()
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
# Clean up the processed file if something went wrong
|
||||
if processed_path.exists():
|
||||
processed_path.unlink(missing_ok=True)
|
||||
raise HTTPException(status_code=500, detail=f"Database error: {str(e)}")
|
||||
|
||||
return {
|
||||
"filename": file.filename,
|
||||
@ -142,37 +204,33 @@ async def upload(request: Request, db = Depends(get_db), uid: str = Form(...), f
|
||||
return {"detail": f"Server error: {type(e).__name__}: {str(e)}"}
|
||||
|
||||
|
||||
def update_public_streams(uid: str, storage_bytes: int, db = Depends(get_db)):
|
||||
def update_public_streams(uid: str, storage_bytes: int, db: Session):
|
||||
"""Update the public streams list in the database with the latest user upload info"""
|
||||
try:
|
||||
from models import PublicStream
|
||||
|
||||
# Get or create the public stream record
|
||||
public_stream = db.get(PublicStream, uid)
|
||||
current_time = datetime.utcnow()
|
||||
|
||||
if public_stream is None:
|
||||
# Create a new record if it doesn't exist
|
||||
public_stream = PublicStream(
|
||||
uid=uid,
|
||||
size=storage_bytes,
|
||||
mtime=int(current_time.timestamp()),
|
||||
created_at=current_time,
|
||||
updated_at=current_time
|
||||
)
|
||||
# Get the user's info
|
||||
user = db.query(User).filter(User.username == uid).first()
|
||||
if not user:
|
||||
print(f"[WARNING] User {uid} not found when updating public streams")
|
||||
return
|
||||
|
||||
# Try to get existing public stream or create new one
|
||||
public_stream = db.query(PublicStream).filter(PublicStream.uid == uid).first()
|
||||
if not public_stream:
|
||||
public_stream = PublicStream(uid=uid)
|
||||
db.add(public_stream)
|
||||
else:
|
||||
# Update existing record
|
||||
public_stream.size = storage_bytes
|
||||
public_stream.mtime = int(current_time.timestamp())
|
||||
public_stream.updated_at = current_time
|
||||
|
||||
# Update the public stream info
|
||||
public_stream.username = user.username
|
||||
public_stream.display_name = user.display_name or user.username
|
||||
public_stream.storage_bytes = storage_bytes
|
||||
public_stream.last_updated = datetime.utcnow()
|
||||
|
||||
db.commit()
|
||||
db.refresh(public_stream)
|
||||
# Don't commit here - let the caller handle the transaction
|
||||
db.flush()
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
# Just log the error and let the caller handle the rollback
|
||||
print(f"[ERROR] Error updating public streams: {e}")
|
||||
import traceback
|
||||
print(f"Error updating public streams in database: {e}")
|
||||
print(traceback.format_exc())
|
||||
raise
|
||||
traceback.print_exc()
|
||||
raise # Re-raise to let the caller handle the error
|
||||
|
||||
Reference in New Issue
Block a user