model-scan-2 / scan_main.py
pengdaqian
init scan
cb34746
raw
history blame
1 kB
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)