server.py

"""Server that will listen for GET and POST requests from the client."""

import time
from typing import List
from fastapi import FastAPI, File, Form, UploadFile, HTTPException
from fastapi.responses import JSONResponse, Response
import asyncio
from common import SERVER_TMP_PATH, SEIZURE_DETECTION_MODEL_PATH
from client_server_interface import FHEServer
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load the server object for seizure detection
FHE_SERVER = FHEServer(model_path=SEIZURE_DETECTION_MODEL_PATH)

def get_server_file_path(name, user_id):
    """Get the correct temporary file path for the server."""
    return SERVER_TMP_PATH / f"{name}_seizure_detection_{user_id}"

# Initialize an instance of FastAPI
app = FastAPI()

@app.get("/")
def root():
    return {"message": "Welcome to Your Seizure Detection FHE Server!"}

@app.post("/send_input")
async def send_input(
    user_id: str = Form(),
    files: List[UploadFile] = File(),
):
    """Send the inputs to the server."""
    try:
        encrypted_image_path = get_server_file_path("encrypted_image", user_id)
        evaluation_key_path = get_server_file_path("evaluation_key", user_id)

        async with encrypted_image_path.open("wb") as encrypted_image, evaluation_key_path.open("wb") as evaluation_key:
            encrypted_image.write(await files[0].read())
            evaluation_key.write(await files[1].read())
        
        return JSONResponse(content={"message": "Input received successfully"}, status_code=200)
    except Exception as e:
        logger.error(f"Error in send_input: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.post("/run_fhe")
async def run_fhe(
    user_id: str = Form(),
):
    """Execute seizure detection on the encrypted input image using FHE."""
    try:
        encrypted_image_path = get_server_file_path("encrypted_image", user_id)
        evaluation_key_path = get_server_file_path("evaluation_key", user_id)

        async with encrypted_image_path.open("rb") as encrypted_image_file, evaluation_key_path.open("rb") as evaluation_key_file:
            encrypted_image = await encrypted_image_file.read()
            evaluation_key = await evaluation_key_file.read()

        async def run_fhe_task():
            start = time.time()
            encrypted_output = FHE_SERVER.run(encrypted_image, evaluation_key)
            fhe_execution_time = round(time.time() - start, 2)

            encrypted_output_path = get_server_file_path("encrypted_output", user_id)
            async with encrypted_output_path.open("wb") as encrypted_output_file:
                await encrypted_output_file.write(encrypted_output)

            return fhe_execution_time

        task = asyncio.create_task(run_fhe_task())
        return JSONResponse(content={"message": "FHE execution started", "task_id": str(id(task))})
    except Exception as e:
        logger.error(f"Error in run_fhe: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.get("/fhe_status/{task_id}")
async def fhe_status(task_id: str):
    """Check the status of an FHE execution task."""
    try:
        for task in asyncio.all_tasks():
            if str(id(task)) == task_id:
                if task.done():
                    try:
                        execution_time = await task
                        return JSONResponse(content={"status": "completed", "execution_time": execution_time})
                    except Exception as e:
                        logger.error(f"Error in task execution: {str(e)}")
                        return JSONResponse(content={"status": "error", "message": str(e)})
                else:
                    return JSONResponse(content={"status": "running"})

        return JSONResponse(content={"status": "not_found"})
    except Exception as e:
        logger.error(f"Error in fhe_status: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.post("/get_output")
async def get_output(
    user_id: str = Form(),
):
    """Retrieve the encrypted output."""
    try:
        encrypted_output_path = get_server_file_path("encrypted_output", user_id)

        async with encrypted_output_path.open("rb") as encrypted_output_file:
            encrypted_output = await encrypted_output_file.read()

        return Response(content=encrypted_output, media_type="application/octet-stream")
    except Exception as e:
        logger.error(f"Error in get_output: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")