rbd-app/src/aeros_project/service.py

import os
from typing import Optional
import json
import httpx
from fastapi import HTTPException, status
from sqlalchemy import Delete, desc, Select, select, func
from sqlalchemy.orm import selectinload

from src.aeros_equipment.service import save_default_equipment
from src.aeros_simulation.service import save_default_simulation_node
from src.auth.service import CurrentUser
from src.config import WINDOWS_AEROS_BASE_URL, CLAMAV_HOST, CLAMAV_PORT
from src.aeros_utils import aeros_post, aeros_file_upload
from src.database.core import DbSession
from src.database.service import search_filter_sort_paginate
from src.utils import sanitize_filename
import clamd
import io

from .model import AerosProject
from .schema import AerosProjectInput, OverhaulScheduleCreate, OverhaulScheduleUpdate
import asyncio
ALLOWED_EXTENSIONS = {".aro"}
MAX_FILE_SIZE = 100 * 1024 * 1024  # 100MB
# client = httpx.AsyncClient(timeout=300.0) # Managed in src.aeros_utils
# We still use a local client for WINDOWS_AEROS_BASE_URL if it's not managed by aeros_utils
client = httpx.AsyncClient(timeout=300.0)


async def import_aro_project(*, db_session: DbSession, aeros_project_in: AerosProjectInput):
    # windows_aeros_base_url = WINDOWS_AEROS_BASE_URL


    file = aeros_project_in.aro_file
    
    # Sanitize and validate filename
    try:
        clean_filename = sanitize_filename(file.filename)
    except ValueError as e:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid filename: {str(e)}"
        )

    # Check if mime type is application/octet-stream 
    if file.content_type != "application/octet-stream":
        raise HTTPException(
            status_code=400,
            detail="Invalid file type. Allowed: application/octet-stream"
        )


    # Get filename
    filename_without_ext = os.path.splitext(clean_filename)[0]

    # Get file extension
    file_ext = os.path.splitext(clean_filename)[1].lower()

    # Validate file extension
    if file_ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"File type not allowed. Allowed: {ALLOWED_EXTENSIONS}"
        )

    print("read file")


    # Read and check file size
    content = await file.read()
    if len(content) > MAX_FILE_SIZE:
        raise HTTPException(
            status_code=400,
            detail="File too large. Max size: 100Mb"
        )

    # ClamAV Scan
    try:
        cd = clamd.ClamdNetworkSocket(CLAMAV_HOST, CLAMAV_PORT)
        scan_result = cd.instream(io.BytesIO(content))
        # Result format: {'stream': ('FOUND', 'Eicar-Test-Signature')} or {'stream': ('OK', None)}
        if scan_result and scan_result.get('stream') and scan_result['stream'][0] == 'FOUND':
            raise HTTPException(
                status_code=400,
                detail=f"Virus detected: {scan_result['stream'][1]}"
            )
    except clamd.ConnectionError:
        raise HTTPException(
            status_code=500,
            detail="Antivirus service unavailable"
        )
    except HTTPException:
        raise
    except Exception as e:
        print(f"ClamAV error: {e}")
        raise HTTPException(
            status_code=500,
            detail=f"Antivirus check failed: {str(e)}"
        )


    # Project name hardcode
    # project_name = "trialapi"

    # Project name
    project_name = filename_without_ext


    ## save File to windows app
    # Output is string of file path, examole
    # # Example response "C/dsad/dsad.aro"
    try:
        # Reset file position since we already read it for size check
        # await file.seek(0)

        # Prepare file for upload
        # files = {
        #     "file": (clean_filename, content, file.content_type or "application/octet-stream")
        # }

        response = await aeros_file_upload(
            "/upload",
            content,
            "file",
            clean_filename
        )

        response.raise_for_status()

        # print("fetch")
        # response = await client.post(
        #     f"{WINDOWS_AEROS_BASE_URL}/upload-file",
        #     files=files
        # )
        # response.raise_for_status()

        # Get the file path from the response
        upload_result = response.json()
        aro_path = upload_result.get("full_path")
        filename = upload_result.get("stored_filename").replace(".aro", "")

        if not aro_path:
            raise HTTPException(
                status_code=500,
                detail="Failed to get file path from upload response"
            )

    except httpx.HTTPStatusError as e:
        raise HTTPException(
            status_code=e.response.status_code,
            detail=f"Upload failed: {e.response.text}"
        )
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    await asyncio.sleep(2)

    # aro_path = r"C:/Users/user/Documents/Aeros/sample_project.aro"

    aeros_project = AerosProject(project_name=filename, aro_file_path=aro_path)

    # find aeros record first, if not found, then create a new one
    stmt = select(AerosProject).order_by(desc(AerosProject.created_at)).limit(1)
    result = await db_session.execute(stmt)
    latest_project = result.scalar_one_or_none()

    # If aeros record found, then update it
    if latest_project:
        latest_project.project_name = filename
        latest_project.aro_file_path = aro_path
    else: # else create new aeros record
        db_session.add(aeros_project)

    await db_session.commit()

    aro_json = json.dumps(aro_path)

    # Update path to AEROS APP
    # Example BODy "C/dsad/dsad.aro"
    try:
        response = await aeros_post(
            "/api/Project/ImportAROFile",
            json_data=aro_json,
            headers={"Content-Type": "application/json"},
        )
        
        # response.raise_for_status() 
        response_json = response.json()
        
        raise Exception(response_json)
        
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    await _initialize_default_project_data(
            db_session=db_session,
            project_name=filename
    )

async def fetch_aro_record(*, db_session: DbSession):
    stmt = select(AerosProject).order_by(desc(AerosProject.updated_at)).limit(1)
    result = await db_session.execute(stmt)
    found_record = result.scalar_one_or_none()

    return found_record

async def get_project(*, db_session: DbSession):
    stmt = select(AerosProject).order_by(desc(AerosProject.updated_at)).limit(1)
    result = await db_session.execute(stmt)
    found_record = result.scalar_one_or_none()

    return found_record

async def _initialize_default_project_data(
    *,
    db_session: DbSession,
    project_name: str
) -> None:
    """
    Initialize default equipment and simulation nodes for a project.

    Args:
        db_session: Database session
        project_name: Name of the project to initialize
    """
    try:
        # Save default equipment
        await save_default_equipment(
            db_session=db_session,
            project_name=project_name
        )

        # # Save default simulation node
        # await save_default_simulation_node(
        #     db_session=db_session,
        #     project_name=project_name
        # )

        await db_session.commit()

    except Exception as e:
        await db_session.rollback()
        raise e


async def reset_project(*, db_session: DbSession):
    project = await fetch_aro_record(db_session=db_session)

    try:
        response = await aeros_post(
            "/api/Project/ImportAROFile",
            data=f'"{project.aro_file_path}"',
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    return project.aro_file_path


async def create(*, db_session: DbSession, overhaul_job_in: OverhaulScheduleCreate):
    # Placeholder for creation logic
    print(f"create overhaul job called: {overhaul_job_in}")
    pass

async def update(*, db_session: DbSession, overhaul_schedule_id: str, overhaul_job_in: OverhaulScheduleUpdate):
    # Placeholder for update logic
    print(f"update overhaul job called for id {overhaul_schedule_id}: {overhaul_job_in}")
    pass

async def delete(*, db_session: DbSession, overhaul_schedule_id: str):
    # Placeholder for delete logic
    print(f"delete overhaul job called for id {overhaul_schedule_id}")
    pass