mirror of
https://github.com/azaion/detections.git
synced 2026-04-22 22:16:31 +00:00
Oleksandr Bezdieniezhnykh
111 lines
2.7 KiB
Python
111 lines
2.7 KiB
Python
import os
|
|
from pathlib import Path
|
|
|
|
from flask import Flask, request
|
|
|
|
app = Flask(__name__)
|
|
|
|
_mode = "normal"
|
|
_first_fail_remaining = False
|
|
_uploads: dict[tuple[str, str], bytes] = {}
|
|
_load_count = 0
|
|
_upload_count = 0
|
|
|
|
|
|
def _models_root() -> Path:
|
|
return Path(os.environ.get("MODELS_ROOT", "/models"))
|
|
|
|
|
|
def _resolve_disk_path(filename: str, folder: str | None) -> Path | None:
|
|
root = _models_root()
|
|
if folder:
|
|
p = root / folder / filename
|
|
else:
|
|
p = root / filename
|
|
if p.is_file():
|
|
return p
|
|
if folder is None:
|
|
alt = root / "models" / filename
|
|
if alt.is_file():
|
|
return alt
|
|
return None
|
|
|
|
|
|
def _should_fail_load() -> bool:
|
|
global _first_fail_remaining
|
|
if _mode == "error":
|
|
return True
|
|
if _mode == "first_fail":
|
|
if _first_fail_remaining:
|
|
_first_fail_remaining = False
|
|
return True
|
|
return False
|
|
return False
|
|
|
|
|
|
@app.route("/load/<path:filename>", methods=["GET", "POST"])
|
|
def load(filename):
|
|
global _load_count
|
|
folder = None
|
|
if request.method == "POST" and request.is_json:
|
|
body = request.get_json(silent=True) or {}
|
|
folder = body.get("folder")
|
|
if _should_fail_load():
|
|
return "", 503
|
|
path = _resolve_disk_path(filename, folder)
|
|
if path is None:
|
|
key = (folder or "", filename)
|
|
data = _uploads.get(key)
|
|
if data is None and folder:
|
|
data = _uploads.get(("", filename))
|
|
if data is None:
|
|
return "", 404
|
|
_load_count += 1
|
|
return data, 200
|
|
_load_count += 1
|
|
return path.read_bytes(), 200
|
|
|
|
|
|
@app.route("/upload/<path:filename>", methods=["POST"])
|
|
def upload(filename):
|
|
global _upload_count
|
|
folder = request.form.get("folder") or ""
|
|
f = request.files.get("data")
|
|
if not f:
|
|
return "", 400
|
|
_uploads[(folder, filename)] = f.read()
|
|
_upload_count += 1
|
|
return "", 200
|
|
|
|
|
|
@app.route("/mock/config", methods=["POST"])
|
|
def mock_config():
|
|
global _mode, _first_fail_remaining
|
|
body = request.get_json(silent=True) or {}
|
|
mode = body.get("mode", "normal")
|
|
if mode not in ("normal", "error", "first_fail"):
|
|
return "", 400
|
|
_mode = mode
|
|
_first_fail_remaining = mode == "first_fail"
|
|
return "", 200
|
|
|
|
|
|
@app.route("/mock/reset", methods=["POST"])
|
|
def mock_reset():
|
|
global _mode, _first_fail_remaining, _uploads, _load_count, _upload_count
|
|
_mode = "normal"
|
|
_first_fail_remaining = False
|
|
_uploads.clear()
|
|
_load_count = 0
|
|
_upload_count = 0
|
|
return "", 200
|
|
|
|
|
|
@app.route("/mock/status", methods=["GET"])
|
|
def mock_status():
|
|
return {
|
|
"mode": _mode,
|
|
"upload_count": _upload_count,
|
|
"load_count": _load_count,
|
|
}
|