scan_main.py

import requests
import uvicorn
from fastapi import BackgroundTasks, FastAPI
import model
from scan import init_clamd, scan_file, clamd_file
from scan_hash import generate_model_hashes

app = FastAPI()


def write_scan_model_result(req: model.ModelScanRequest):
    ret = {}
    if 'Scan' in req.tasks:
        detail = scan_file(req.fileUrl)
        clamd_detail = clamd_file(req.fileUrl, clamd_exec)
        ret |= detail
        ret |= clamd_detail
    if 'Hash' in req.tasks:
        ret["hashes"] = generate_model_hashes(req.fileUrl)

    try:
        requests.post(req.callbackUrl, json=ret)
    except Exception as ex:
        print(ex)


@app.post("/model-scan")
async def model_scan_handler(req: model.ModelScanRequest, background_tasks: BackgroundTasks):
    background_tasks.add_task(write_scan_model_result, req)
    return model.ModelScanResponse(ok=True, error="")


if __name__ == "__main__":
    global clamd_exec
    clamd_exec = init_clamd()
    uvicorn.run(app, host="0.0.0.0", port=7860)