rbox/mywebdav/routers/shares.py

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