149 lines
5.4 KiB
Python
149 lines
5.4 KiB
Python
import hashlib
|
|
import struct
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Generator
|
|
|
|
from argon2 import PasswordHasher
|
|
from argon2.exceptions import VerificationError, VerifyMismatchError
|
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
from fastapi import APIRouter, Depends, HTTPException, Request
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
from sqlalchemy.orm import Session
|
|
|
|
from config import settings
|
|
from database import Share, SecurityEvent, get_db
|
|
|
|
router = APIRouter()
|
|
_ph = PasswordHasher()
|
|
|
|
# ── Helpers ────────────────────────────────────────────────────────────────
|
|
|
|
def _hash_ip(ip: str) -> str:
|
|
return hashlib.sha256(ip.encode()).hexdigest()
|
|
|
|
def _client_ip(request: Request) -> str:
|
|
fwd = request.headers.get('X-Forwarded-For')
|
|
if fwd:
|
|
return fwd.split(',')[0].strip()
|
|
return request.client.host if request.client else 'unknown'
|
|
|
|
def _live_share(share_id: str, db: Session) -> Share:
|
|
share = db.query(Share).filter(Share.id == share_id).first()
|
|
if not share:
|
|
raise HTTPException(404, 'Share not found')
|
|
if datetime.utcnow() > share.expires_at:
|
|
raise HTTPException(404, 'Share not found')
|
|
if share.burn_after_read and share.downloaded:
|
|
raise HTTPException(404, 'Share not found')
|
|
return share
|
|
|
|
def _decrypt_stream(path: Path, key_hex: str) -> Generator[bytes, None, None]:
|
|
"""Yield decrypted plaintext chunks from a chunked AES-256-GCM file."""
|
|
aesgcm = AESGCM(bytes.fromhex(key_hex))
|
|
with open(path, 'rb') as f:
|
|
chunk_count = struct.unpack('>I', f.read(4))[0]
|
|
for _ in range(chunk_count):
|
|
length = struct.unpack('>I', f.read(4))[0]
|
|
payload = f.read(length)
|
|
nonce, ct = payload[:12], payload[12:]
|
|
yield aesgcm.decrypt(nonce, ct, None)
|
|
|
|
# ── Routes ─────────────────────────────────────────────────────────────────
|
|
|
|
@router.get('/share/{share_id}/meta')
|
|
async def share_meta(share_id: str, db: Session = Depends(get_db)):
|
|
share = _live_share(share_id, db)
|
|
|
|
locked, lock_remaining = False, 0
|
|
if share.locked_until and datetime.utcnow() < share.locked_until:
|
|
locked = True
|
|
lock_remaining = int((share.locked_until - datetime.utcnow()).total_seconds())
|
|
|
|
return {
|
|
'share_type' : share.share_type,
|
|
'original_filename': share.original_filename,
|
|
'mimetype' : share.mimetype,
|
|
'filesize' : share.filesize,
|
|
'expires_at' : share.expires_at.isoformat() + 'Z',
|
|
'burn_after_read' : share.burn_after_read,
|
|
'has_password' : share.password_hash is not None,
|
|
'locked' : locked,
|
|
'lock_remaining' : lock_remaining,
|
|
}
|
|
|
|
|
|
class PasswordBody(BaseModel):
|
|
password: str = ''
|
|
|
|
|
|
@router.post('/share/{share_id}/download')
|
|
async def download_share(
|
|
share_id: str,
|
|
body: PasswordBody,
|
|
request: Request,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
share = _live_share(share_id, db)
|
|
now = datetime.utcnow()
|
|
ip_hash = _hash_ip(_client_ip(request))
|
|
|
|
# Lockout check
|
|
if share.locked_until and now < share.locked_until:
|
|
secs = int((share.locked_until - now).total_seconds())
|
|
raise HTTPException(429, f'Too many failed attempts. Try again in {secs}s.')
|
|
|
|
# Password check
|
|
if share.password_hash:
|
|
if not body.password:
|
|
raise HTTPException(401, 'Password required')
|
|
try:
|
|
_ph.verify(share.password_hash, body.password)
|
|
except (VerifyMismatchError, VerificationError):
|
|
share.failed_attempts = (share.failed_attempts or 0) + 1
|
|
db.add(SecurityEvent(
|
|
event_type='failed_attempt', share_id=share_id,
|
|
ip_hash=ip_hash, details=f'attempt #{share.failed_attempts}',
|
|
))
|
|
if share.failed_attempts >= settings.max_failed_attempts:
|
|
share.locked_until = now + timedelta(minutes=settings.lockout_minutes)
|
|
db.add(SecurityEvent(
|
|
event_type='share_locked', share_id=share_id, ip_hash=ip_hash,
|
|
details=f'locked {settings.lockout_minutes}m after {share.failed_attempts} attempts',
|
|
))
|
|
db.commit()
|
|
raise HTTPException(401, 'Incorrect password')
|
|
|
|
share.failed_attempts = 0
|
|
share.locked_until = None
|
|
|
|
blob_path = Path(settings.storage_path) / share.storage_path
|
|
if not blob_path.exists():
|
|
raise HTTPException(404, 'Share not found')
|
|
|
|
share.downloaded = True
|
|
db.add(SecurityEvent(event_type='share_accessed', share_id=share_id))
|
|
db.commit()
|
|
|
|
key_hex = share.encryption_key
|
|
burn = share.burn_after_read
|
|
|
|
def _stream():
|
|
try:
|
|
yield from _decrypt_stream(blob_path, key_hex)
|
|
finally:
|
|
if burn:
|
|
blob_path.unlink(missing_ok=True)
|
|
|
|
headers = {}
|
|
if share.original_filename:
|
|
safe = share.original_filename.encode('ascii', 'replace').decode()
|
|
headers['Content-Disposition'] = f'attachment; filename="{safe}"'
|
|
|
|
return StreamingResponse(
|
|
_stream(),
|
|
media_type=share.mimetype or 'application/octet-stream',
|
|
headers=headers,
|
|
)
|