rbox/shares.py at master

 from fastapi import APIRouter, Depends, HTTPException, status
 from fastapi.responses import StreamingResponse
 from typing import Optional, List
 import secrets
 from datetime import datetime
 from ..auth import get_current_user
 from ..models import User, File, Folder, Share
 from ..schemas import ShareCreate, ShareOut, FileOut, FolderOut
 from ..auth import get_password_hash, verify_password
 from ..storage import storage_manager
 from ..mail import send_email
 from ..settings import settings
 router = APIRouter(
     prefix="/shares",
     tags=["shares"],
 )
 @router.post("/", response_model=ShareOut, status_code=status.HTTP_201_CREATED)
 async def create_share_link(
     share_in: ShareCreate, current_user: User = Depends(get_current_user)
 ):
     if not share_in.file_id and not share_in.folder_id:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail="Either file_id or folder_id must be provided",
         )
     if share_in.file_id and share_in.folder_id:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail="Cannot share both a file and a folder simultaneously",
         )
     file = None
     folder = None
     if share_in.file_id:
         file = await File.get_or_none(
             id=share_in.file_id, owner=current_user, is_deleted=False
         )
         if not file:
             raise HTTPException(
                 status_code=status.HTTP_404_NOT_FOUND,
                 detail="File not found or does not belong to you",
             )
     if share_in.folder_id:
         folder = await Folder.get_or_none(
             id=share_in.folder_id, owner=current_user, is_deleted=False
         )
         if not folder:
             raise HTTPException(
                 status_code=status.HTTP_404_NOT_FOUND,
                 detail="Folder not found or does not belong to you",
             )
     token = secrets.token_urlsafe(16)
     hashed_password = None
     password_protected = False
     if share_in.password:
         hashed_password = get_password_hash(share_in.password)
         password_protected = True
     share = await Share.create(
         token=token,
         file=file,
         folder=folder,
         owner=current_user,
         expires_at=share_in.expires_at,
         password_protected=password_protected,
         hashed_password=hashed_password,
         permission_level=share_in.permission_level,
     )
     if share_in.invite_email:
         share_url = f"https://{settings.DOMAIN_NAME}/share/{token}"
         item_type = "file" if file else "folder"
         item_name = file.name if file else folder.name
         expiry_text = (
             f" until {share_in.expires_at.strftime('%Y-%m-%d %H:%M')}"
             if share_in.expires_at
             else ""
         )
         password_text = (
             f"\n\nPassword: {share_in.password}" if share_in.password else ""
         )
         email_body = f"""Hello,
 {current_user.username} has shared a {item_type} with you: {item_name}
 Permission level: {share_in.permission_level}
 Access link: {share_url}{password_text}
 This link is valid{expiry_text}.
 --
 MyWebdav File Sharing Service"""
         email_html = f"""
 <html>
 <body>
     <p>Hello,</p>
     <p><strong>{current_user.username}</strong> has shared a {item_type} with you: <strong>{item_name}</strong></p>
     <p><strong>Permission level:</strong> {share_in.permission_level}</p>
     <p><a href="{share_url}" style="background-color: #4CAF50; color: white; padding: 10px 20px; text-decoration: none; border-radius: 4px; display: inline-block;">Access {item_type.capitalize()}</a></p>
     {f'<p><strong>Password:</strong> {share_in.password}</p>' if share_in.password else ''}
     <p><small>This link is valid{expiry_text}.</small></p>
     <hr>
     <p><small>MyWebdav File Sharing Service</small></p>
 </body>
 </html>
 """
         try:
             await send_email(
                 to_email=share_in.invite_email,
                 subject=f"{current_user.username} shared {item_name} with you",
                 body=email_body,
                 html=email_html,
             )
         except Exception as e:
             print(f"Failed to send invitation email: {e}")
     return await ShareOut.from_tortoise_orm(share)
 @router.get("/my", response_model=List[ShareOut])
 async def list_my_shares(current_user: User = Depends(get_current_user)):
     shares = await Share.filter(owner=current_user).order_by("-created_at")
     return [await ShareOut.from_tortoise_orm(share) for share in shares]
 @router.get("/{share_token}", response_model=ShareOut)
 async def get_share_link_info(share_token: str):
     share = await Share.get_or_none(token=share_token)
     if not share:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND, detail="Share link not found"
         )
     if share.expires_at and share.expires_at < datetime.utcnow():
         raise HTTPException(
             status_code=status.HTTP_410_GONE, detail="Share link has expired"
         )
     # Increment access count
     share.access_count += 1
     await share.save()
     return await ShareOut.from_tortoise_orm(share)
 @router.put("/{share_id}", response_model=ShareOut)
 async def update_share(
     share_id: int, share_in: ShareCreate, current_user: User = Depends(get_current_user)
 ):
     share = await Share.get_or_none(id=share_id, owner=current_user)
     if not share:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="Share link not found or does not belong to you",
         )
     if share_in.expires_at is not None:
         share.expires_at = share_in.expires_at
     if share_in.password is not None:
         share.hashed_password = get_password_hash(share_in.password)
         share.password_protected = True
     elif share_in.password == "":  # Allow clearing password
         share.hashed_password = None
         share.password_protected = False
     if share_in.permission_level is not None:
         share.permission_level = share_in.permission_level
     await share.save()
     return await ShareOut.from_tortoise_orm(share)
 @router.post("/{share_token}/access")
 async def access_shared_content(share_token: str, password: Optional[str] = None):
     share = await Share.get_or_none(token=share_token)
     if not share:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND, detail="Share link not found"
         )
     if share.expires_at and share.expires_at < datetime.utcnow():
         raise HTTPException(
             status_code=status.HTTP_410_GONE, detail="Share link has expired"
         )
     if share.password_protected:
         if not password or not verify_password(password, share.hashed_password):
             raise HTTPException(
                 status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect password"
             )
     result = {"message": "Access granted", "permission_level": share.permission_level}
     if share.file_id:
         file = await File.get_or_none(id=share.file_id, is_deleted=False)
         if not file:
             raise HTTPException(
                 status_code=status.HTTP_404_NOT_FOUND, detail="File not found"
             )
         result["file"] = await FileOut.from_tortoise_orm(file)
         result["type"] = "file"
     elif share.folder_id:
         folder = await Folder.get_or_none(id=share.folder_id, is_deleted=False)
         if not folder:
             raise HTTPException(
                 status_code=status.HTTP_404_NOT_FOUND, detail="Folder not found"
             )
         result["folder"] = await FolderOut.from_tortoise_orm(folder)
         result["type"] = "folder"
         files = await File.filter(parent=folder, is_deleted=False)
         subfolders = await Folder.filter(parent=folder, is_deleted=False)
         result["files"] = [await FileOut.from_tortoise_orm(f) for f in files]
         result["folders"] = [await FolderOut.from_tortoise_orm(f) for f in subfolders]
     return result
 @router.get("/{share_token}/download")
 async def download_shared_file(share_token: str, password: Optional[str] = None):
     share = await Share.get_or_none(token=share_token)
     if not share:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND, detail="Share link not found"
         )
     if share.expires_at and share.expires_at < datetime.utcnow():
         raise HTTPException(
             status_code=status.HTTP_410_GONE, detail="Share link has expired"
         )
     if share.password_protected:
         if not password or not verify_password(password, share.hashed_password):
             raise HTTPException(
                 status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect password"
             )
     if not share.file_id:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail="This share is not for a file",
         )
     file = await File.get_or_none(id=share.file_id, is_deleted=False)
     if not file:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND, detail="File not found"
         )
     owner = await User.get(id=file.owner_id)
     try:
         async def file_iterator():
             async for chunk in storage_manager.get_file(owner.id, file.path):
                 yield chunk
         return StreamingResponse(
             content=file_iterator(),
             media_type=file.mime_type,
             headers={"Content-Disposition": f'attachment; filename="{file.name}"'},
         )
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="File not found in storage")
 @router.delete("/{share_id}", status_code=status.HTTP_204_NO_CONTENT)
 async def delete_share_link(
     share_id: int, current_user: User = Depends(get_current_user)
 ):
     share = await Share.get_or_none(id=share_id, owner=current_user)
     if not share:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="Share link not found or does not belong to you",
         )
     await share.delete()
     return