personal_development/ServerManagerWebApp/backend/app/routers/samba.py

from fastapi import APIRouter, UploadFile, File, Request, Response
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from ..samba_manager import SambaManager
from ..config import settings
import os
import io
import mimetypes
import zipfile

class RenameRequest(BaseModel):
    old_path: str
    new_path: str

router = APIRouter()


@router.get("/list/{share_name}")
def list_files(share_name: str, path: str = "/"):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    files = manager.list_path(path)
    manager.close()
    return files


@router.delete("/files/{share_name}")
def delete_file(share_name: str, path: str):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    result = manager.delete_file(path)
    manager.close()
    return result


@router.post("/files/upload/{share_name}")
def upload_file(share_name: str, path: str, file: UploadFile = File(...)):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    # The path in storeFile is service_name + path
    remote_path = os.path.join(path, file.filename)
    result = manager.upload_file(remote_path, file.file)
    manager.close()
    return result


@router.get("/files/download/{share_name}")
def download_file(share_name: str, path: str):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    file_obj = io.BytesIO()
    manager.download_file(path, file_obj)
    manager.close()
    file_obj.seek(0)
    return StreamingResponse(file_obj, media_type="application/octet-stream")


@router.get("/view/{share_name}")
def view_file(share_name: str, path: str, request: Request):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    
    try:
        attributes = manager.get_attributes(path)
        if isinstance(attributes, dict) and "error" in attributes:
            return Response(status_code=404)
        file_size = attributes.file_size

        range_header = request.headers.get('Range')
        
        media_type, _ = mimetypes.guess_type(path)
        if media_type is None:
            media_type = "application/octet-stream"

        if range_header:
            start_str, end_str = range_header.replace('bytes=', '').split('-')
            start = int(start_str)
            end = int(end_str) if end_str else file_size - 1

            if start >= file_size:
                return Response(status_code=416)

            length = end - start + 1
            file_obj = io.BytesIO()
            manager.download_file_range(path, file_obj, offset=start, max_length=length)
            file_obj.seek(0)
            
            headers = {
                'Content-Range': f'bytes {start}-{end}/{file_size}',
                'Accept-Ranges': 'bytes',
                'Content-Length': str(length),
            }
            return Response(file_obj.read(), status_code=206, headers=headers, media_type=media_type)
        else:
            file_obj = io.BytesIO()
            manager.download_file(path, file_obj)
            file_obj.seek(0)
            return StreamingResponse(file_obj, media_type=media_type, headers={"Accept-Ranges": "bytes", "Content-Length": str(file_size)})
    finally:
        manager.close()


@router.get("/cbz/list/{share_name}")
def list_cbz_contents(share_name: str, path: str):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    file_obj = io.BytesIO()
    manager.download_file(path, file_obj)
    manager.close()
    file_obj.seek(0)

    image_files = []
    with zipfile.ZipFile(file_obj, 'r') as zip_ref:
        for file_info in zip_ref.infolist():
            if file_info.filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp')):
                image_files.append(file_info.filename)
    
    return sorted(image_files)

@router.get("/cbz/image/{share_name}")
def get_cbz_image(share_name: str, path: str, image_name: str):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    file_obj = io.BytesIO()
    manager.download_file(path, file_obj)
    manager.close()
    file_obj.seek(0)

    with zipfile.ZipFile(file_obj, 'r') as zip_ref:
        with zip_ref.open(image_name) as image_file:
            image_data = io.BytesIO(image_file.read())
    
    media_type, _ = mimetypes.guess_type(image_name)
    if media_type is None:
        media_type = "application/octet-stream"
        
    return StreamingResponse(image_data, media_type=media_type)


@router.put("/files/rename/{share_name}")
def rename_file(share_name: str, request: RenameRequest):
    manager = SambaManager(
        settings.samba_server_ip,
        share_name,
        username=settings.samba_username,
        password=settings.samba_password,
    )
    result = manager.rename_file(request.old_path, request.new_path)
    manager.close()
    return result