84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
"""
|
|
Blueprint: API DrHider.
|
|
|
|
Three endpoints:
|
|
- POST /api/upload — upload one file -> {session_id}
|
|
- POST /api/process/{sid} — process all session files -> {status:"done"}
|
|
- GET /api/download/{sid} — download ZIP
|
|
"""
|
|
|
|
import io
|
|
import zipfile
|
|
import traceback
|
|
from flask import Blueprint, request, send_file, jsonify
|
|
|
|
from drhider import obfuscate_files, LLMClient
|
|
from drhider.builder import build_zip, build_mapping_csv
|
|
from session import create_session, add_file, get_files, store_result, get_result, cleanup, file_count
|
|
|
|
api_bp = Blueprint("api", __name__, url_prefix="/api")
|
|
|
|
|
|
@api_bp.route("/upload", methods=["POST"])
|
|
def upload():
|
|
"""Upload one file to session."""
|
|
sid = request.form.get("session", "")
|
|
if not sid:
|
|
sid = create_session()
|
|
uploaded = request.files.getlist("files")
|
|
if not uploaded:
|
|
return jsonify({"ok": False, "error": "No file"}), 400
|
|
f = uploaded[0]
|
|
if not f.filename:
|
|
return jsonify({"ok": False, "error": "No filename"}), 400
|
|
ok = add_file(sid, f.filename, f.read())
|
|
if not ok:
|
|
return jsonify({"ok": False, "error": "Session not found"}), 404
|
|
return jsonify({"ok": True, "session": sid, "count": file_count(sid)})
|
|
|
|
|
|
@api_bp.route("/process/<sid>", methods=["POST"])
|
|
def process(sid):
|
|
"""Process all session files one by one -> ZIP."""
|
|
files = get_files(sid)
|
|
if files is None:
|
|
return jsonify({"ok": False, "error": "Session not found"}), 404
|
|
if not files:
|
|
return jsonify({"ok": False, "error": "No files"}), 400
|
|
try:
|
|
llm = LLMClient()
|
|
all_results = []
|
|
all_mapping = {}
|
|
for fname, content in files:
|
|
zip_data, _csv = obfuscate_files([(fname, content, "")], llm_client=llm)
|
|
with zipfile.ZipFile(io.BytesIO(zip_data)) as zf:
|
|
for name in zf.namelist():
|
|
if name == "mapping.csv":
|
|
csv_text = zf.read(name).decode("utf-8")
|
|
lines = csv_text.strip().split("\n")
|
|
if len(lines) > 1:
|
|
for line in lines[1:]:
|
|
parts = line.split(",", 2)
|
|
if len(parts) >= 3:
|
|
all_mapping[parts[1]] = parts[2]
|
|
elif not name.endswith("/"):
|
|
all_results.append((name, zf.read(name)))
|
|
csv_str = build_mapping_csv(all_mapping) if all_mapping else ""
|
|
final_zip = build_zip(all_results, mapping_csv=csv_str)
|
|
store_result(sid, final_zip)
|
|
return jsonify({"ok": True, "status": "done", "files": len(all_results)})
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
return jsonify({"ok": False, "error": str(e)}), 500
|
|
|
|
|
|
@api_bp.route("/download/<sid>", methods=["GET"])
|
|
def download(sid):
|
|
"""Download result and cleanup session."""
|
|
zip_data = get_result(sid)
|
|
if zip_data is None:
|
|
return jsonify({"ok": False, "error": "Not found"}), 404
|
|
cleanup(sid)
|
|
return send_file(io.BytesIO(zip_data), mimetype="application/zip",
|
|
as_attachment=True, download_name="drhider_output.zip")
|