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)