Refactor: move API endpoints to /api prefix, use flask-openapi3 (#1068)

* Refactor: flask-openapi3

* Use RootModel

* web: Set baseURL to /api

* Fix Makefile

* Fix request declarations

* Separate schema module

* Screenshot fixes

* Add API docs to navbar

* Fix linting

Makefile

@@ -1,7 +1,7 @@
 rwildcard = $(foreach d,$(wildcard $1*),\
                 $(call rwildcard,$d/,$2) \
                 $(filter $(subst *,%,$2),$d))
-WEB_SOURCE_FILES := $(call rwildcard,drakrun/web/frontend/src/,*.js) $(call rwildcard,drakrun/web/frontend/src/,*.css)
+WEB_SOURCE_FILES := $(call rwildcard,drakrun/web/frontend/src/,*.js) $(call rwildcard,drakrun/web/frontend/src/,*.jsx) $(call rwildcard,drakrun/web/frontend/src/,*.css)
 PYTHON_SOURCE_FILES := $(call rwildcard,drakrun/,*.py) drakrun/data pyproject.toml MANIFEST.in requirements.txt setup.py
 
 .PHONY: all

drakrun/web/api.py

@@ -0,0 +1,260 @@
+import hashlib
+import json
+import logging
+import shutil
+import uuid
+from tempfile import NamedTemporaryFile
+from zipfile import ZIP_DEFLATED, ZipFile
+
+import magic
+import orjson
+from flask import Response, jsonify, request, send_file
+from flask_openapi3 import APIBlueprint
+from rq.exceptions import NoSuchJobError
+from rq.job import Job, JobStatus
+
+from drakrun.analyzer.analysis_options import AnalysisOptions
+from drakrun.analyzer.file_metadata import FileMetadata
+from drakrun.analyzer.postprocessing.indexer import scattered_read_file
+from drakrun.analyzer.worker import (
+    analysis_job_to_status_dict,
+    enqueue_analysis,
+    get_redis_connection,
+)
+from drakrun.lib.config import load_config
+from drakrun.lib.paths import ANALYSES_DIR
+from drakrun.web.analysis import get_analysis_data
+from drakrun.web.analysis_list import add_analysis_to_recent, get_recent_analysis_list
+from drakrun.web.schema import (
+    AnalysisListResponse,
+    AnalysisRequestPath,
+    AnalysisResponse,
+    APIErrorResponse,
+    LogsRequestPath,
+    ProcessedRequestPath,
+    ProcessInfoRequestPath,
+    ProcessLogsRequestPath,
+    ScreenshotRequestPath,
+    UploadAnalysisResponse,
+    UploadFileForm,
+)
+
+api = APIBlueprint("api", __name__, url_prefix="/api")
+
+config = load_config()
+redis = get_redis_connection(config.redis)
+logger = logging.getLogger(__name__)
+
+
+@api.post("/upload", responses={200: UploadAnalysisResponse})
+def upload_sample(form: UploadFileForm):
+    request_file = form.file
+    job_id = str(uuid.uuid4())
+
+    analysis_path = ANALYSES_DIR / job_id
+    analysis_path.mkdir()
+
+    sample_path = analysis_path / "sample"
+    sample_sha256 = hashlib.sha256()
+    try:
+        with sample_path.open("wb") as f:
+            for chunk in iter(lambda: request_file.read(32 * 4096), b""):
+                f.write(chunk)
+                sample_sha256.update(chunk)
+        sample_magic = magic.from_file(sample_path.as_posix())
+
+        timeout = form.timeout
+        if not timeout:
+            timeout = config.default_timeout
+        filename = form.file_name
+        if not filename:
+            filename = request_file.filename
+        start_command = form.start_command
+        plugins = form.plugins
+        if plugins:
+            plugins = json.loads(plugins)
+
+        file_metadata = FileMetadata(
+            name=filename,
+            type=sample_magic,
+            sha256=sample_sha256.hexdigest(),
+        )
+        analysis_options = AnalysisOptions(
+            config=config,
+            sample_path=sample_path,
+            target_filename=filename,
+            start_command=start_command,
+            plugins=plugins,
+            timeout=timeout,
+            job_timeout_leeway=config.drakrun.job_timeout_leeway,
+        )
+        enqueue_analysis(
+            job_id=job_id,
+            file_metadata=file_metadata,
+            options=analysis_options,
+            connection=redis,
+        )
+        add_analysis_to_recent(connection=redis, analysis_id=job_id)
+        return jsonify({"task_uid": job_id})
+    except Exception:
+        shutil.rmtree(analysis_path)
+        raise
+
+
+@api.get("/list", responses={200: AnalysisListResponse})
+def list_analyses():
+    analysis_list = get_recent_analysis_list(redis)
+    return jsonify([analysis_job_to_status_dict(job) for job in analysis_list])
+
+
+@api.get("/status/<task_uid>", responses={200: AnalysisResponse, 404: APIErrorResponse})
+def status(path: AnalysisRequestPath):
+    task_uid = path.task_uid
+    try:
+        job = Job.fetch(task_uid, connection=redis)
+    except NoSuchJobError:
+        job = None
+
+    if job is not None and job.get_status() not in [
+        JobStatus.FINISHED,
+        JobStatus.FAILED,
+    ]:
+        return jsonify(analysis_job_to_status_dict(job))
+
+    analysis = get_analysis_data(task_uid)
+    metadata = analysis.get_metadata()
+    # Handling old tasks, to be removed in future
+    if "id" not in metadata:
+        metadata = {"id": task_uid, **metadata}
+    if metadata is None:
+        return jsonify({"error": "Job not found"}), 404
+    else:
+        return jsonify(metadata)
+
+
+@api.get("/processed/<task_uid>/<which>")
+def processed(path: ProcessedRequestPath):
+    task_uid = path.task_uid
+    which = path.which
+    analysis = get_analysis_data(task_uid)
+    path = analysis.get_processed(which)
+    if not path.exists():
+        return dict(error="Data not found"), 404
+    return send_file(path, mimetype="application/json")
+
+
+@api.get("/logs/<task_uid>/<log_type>")
+def logs(path: LogsRequestPath):
+    task_uid = path.task_uid
+    log_type = path.log_type
+    analysis = get_analysis_data(task_uid)
+    path = analysis.get_log(log_type)
+    if not path.exists():
+        return dict(error="Data not found"), 404
+    return send_file(path, mimetype="text/plain")
+
+
+@api.get("/process_info/<task_uid>/<seqid>")
+def process_info(path: ProcessInfoRequestPath):
+    task_uid = path.task_uid
+    seqid = path.seqid
+    analysis = get_analysis_data(task_uid)
+    process_info = analysis.get_process_info(seqid)
+    if process_info is None:
+        return dict(error="Data not found"), 404
+    return jsonify(process_info)
+
+
+@api.get("/logs/<task_uid>/<log_type>/process/<seqid>")
+def process_logs(path: ProcessLogsRequestPath):
+    task_uid = path.task_uid
+    log_type = path.log_type
+    seqid = path.seqid
+    analysis = get_analysis_data(task_uid)
+    index_path = analysis.get_log_index(f"{log_type}.{seqid}.json")
+    if not index_path.exists():
+        return dict(error="Data not found"), 404
+    index = orjson.loads(index_path.read_text())
+    blocks = index["blocks"]
+    filter_values = request.args.getlist("filter[]")
+    if filter_values:
+        filter_indices = [
+            index["values"].index(filter_value)
+            for filter_value in filter_values
+            if filter_value in index["values"]
+        ]
+        blocks = [
+            block
+            for (block, mapping) in zip(index["blocks"], index["mapping"])
+            if mapping in filter_indices
+        ]
+    log_path = analysis.get_log(log_type)
+    if request.range:
+        if len(request.range.ranges) > 1:
+            return dict(error="Multiple ranges unsupported"), 400
+        range_start, range_stop = request.range.ranges[0]
+        skip = range_start
+        length = (range_stop - range_start + 1) if range_stop is not None else None
+    else:
+        skip, length = 0, None
+    scattered_read = scattered_read_file(log_path, blocks, skip=skip, length=length)
+    return Response(b"".join(scattered_read), mimetype="text/plain")
+
+
+@api.get("/pcap_dump/<task_uid>")
+def pcap_dump(path: AnalysisRequestPath):
+    """
+    Return archive containing dump.pcap along with extracted tls sessions
+    keys in format acceptable by wireshark.
+    """
+    task_uid = path.task_uid
+    analysis = get_analysis_data(task_uid)
+    path = analysis.get_pcap_dump()
+    if not path.exists():
+        return dict(error="Data not found"), 404
+    with NamedTemporaryFile() as f_archive:
+        with ZipFile(f_archive, "w", ZIP_DEFLATED) as archive:
+            archive.write(path, "dump.pcap")
+            path = analysis.get_wireshark_key_file()
+            if path.exists():
+                archive.write(path, "dump.keys")
+        f_archive.seek(0)
+        return send_file(f_archive.name, mimetype="application/zip")
+
+
+@api.get("/dumps/<task_uid>")
+def dumps(path: AnalysisRequestPath):
+    task_uid = path.task_uid
+    analysis = get_analysis_data(task_uid)
+    path = analysis.get_dumps()
+    if not path.exists():
+        return dict(error="Data not found"), 404
+    return send_file(path, mimetype="application/zip")
+
+
+@api.get("/logs/<task_uid>")
+def list_logs(path: AnalysisRequestPath):
+    task_uid = path.task_uid
+    analysis = get_analysis_data(task_uid)
+    return jsonify(list(analysis.list_logs()))
+
+
+@api.get("/graph/<task_uid>")
+def graph(path: AnalysisRequestPath):
+    task_uid = path.task_uid
+    analysis = get_analysis_data(task_uid)
+    path = analysis.get_graph()
+    if not path.exists():
+        return dict(error="Data not found"), 404
+    return send_file(path, mimetype="text/plain")
+
+
+@api.get("/screenshot/<task_uid>/<which>")
+def screenshot(path: ScreenshotRequestPath):
+    task_uid = path.task_uid
+    which = path.which
+    analysis = get_analysis_data(task_uid)
+    path = analysis.get_screenshot(which)
+    if not path.exists():
+        return dict(error="Data not found"), 404
+    return send_file(path, mimetype="image/png")

drakrun/web/app.py

@@ -1,42 +1,25 @@
-import hashlib
-import json
 import os
-import shutil
-import uuid
-from tempfile import NamedTemporaryFile
-from zipfile import ZIP_DEFLATED, ZipFile
 
-import magic
 import rq_dashboard
-from flask import Flask, Response, jsonify, request, send_file
+from flask import jsonify, send_file
-from orjson import orjson
+from flask_openapi3 import Info, OpenAPI
-from rq.exceptions import NoSuchJobError
-from rq.job import Job, JobStatus
 
-from drakrun.analyzer.analysis_options import AnalysisOptions
-from drakrun.analyzer.file_metadata import FileMetadata
-from drakrun.analyzer.postprocessing.indexer import scattered_read_file
-from drakrun.analyzer.worker import (
-    analysis_job_to_status_dict,
-    enqueue_analysis,
-    get_redis_connection,
-)
 from drakrun.lib.config import load_config
-from drakrun.lib.paths import ANALYSES_DIR
+from drakrun.version import __version__
 
-from .analysis import get_analysis_data
+from .api import api
-from .analysis_list import add_analysis_to_recent, get_recent_analysis_list
 
-app = Flask(__name__, static_folder="frontend/dist/assets")
+info = Info(title="Drakvuf Sandbox", version=__version__)
-drakrun_conf = load_config()
+app = OpenAPI(__name__, info=info, static_folder="frontend/dist/assets")
-redis = get_redis_connection(drakrun_conf.redis)
+config = load_config()
 app.config.update(
     {
-        "RQ_DASHBOARD_REDIS_URL": drakrun_conf.redis.make_url(),
+        "RQ_DASHBOARD_REDIS_URL": config.redis.make_url(),
     }
 )
 rq_dashboard.web.setup_rq_connection(app)
 app.register_blueprint(rq_dashboard.blueprint, url_prefix="/rq")
+app.register_api(api)
 
 
 @app.errorhandler(404)
@@ -53,216 +36,6 @@ if os.environ.get("DRAKRUN_CORS_ALL"):
         return response
 
 
-@app.route("/upload", methods=["POST"])
-def upload_sample():
-    if "file" not in request.files:
-        return jsonify(error="No file part"), 400
-
-    request_file = request.files["file"]
-    job_id = str(uuid.uuid4())
-
-    analysis_path = ANALYSES_DIR / job_id
-    analysis_path.mkdir()
-
-    sample_path = analysis_path / "sample"
-    sample_sha256 = hashlib.sha256()
-    try:
-        with sample_path.open("wb") as f:
-            for chunk in iter(lambda: request_file.read(32 * 4096), b""):
-                f.write(chunk)
-                sample_sha256.update(chunk)
-        sample_magic = magic.from_file(sample_path.as_posix())
-
-        timeout = request.form.get("timeout")
-        if not timeout:
-            timeout = drakrun_conf.default_timeout
-        filename = request.form.get("file_name")
-        if not filename:
-            filename = request_file.filename
-        start_command = request.form.get("start_command")
-        plugins = request.form.get("plugins")
-        if plugins:
-            plugins = json.loads(plugins)
-
-        file_metadata = FileMetadata(
-            name=filename,
-            type=sample_magic,
-            sha256=sample_sha256.hexdigest(),
-        )
-        analysis_options = AnalysisOptions(
-            config=drakrun_conf,
-            sample_path=sample_path,
-            target_filename=filename,
-            start_command=start_command,
-            plugins=plugins,
-            timeout=timeout,
-            job_timeout_leeway=drakrun_conf.drakrun.job_timeout_leeway,
-        )
-        enqueue_analysis(
-            job_id=job_id,
-            file_metadata=file_metadata,
-            options=analysis_options,
-            connection=redis,
-        )
-        add_analysis_to_recent(connection=redis, analysis_id=job_id)
-        return jsonify({"task_uid": job_id})
-    except Exception:
-        shutil.rmtree(analysis_path)
-        raise
-
-
-@app.route("/list")
-def list_analyses():
-    analysis_list = get_recent_analysis_list(redis)
-    return jsonify([analysis_job_to_status_dict(job) for job in analysis_list])
-
-
-@app.route("/status/<task_uid>")
-def status(task_uid):
-    try:
-        job = Job.fetch(task_uid, connection=redis)
-    except NoSuchJobError:
-        job = None
-
-    if job is not None and job.get_status() not in [
-        JobStatus.FINISHED,
-        JobStatus.FAILED,
-    ]:
-        return jsonify(analysis_job_to_status_dict(job))
-
-    analysis = get_analysis_data(task_uid)
-    metadata = analysis.get_metadata()
-    # Handling old tasks, to be removed in future
-    if "id" not in metadata:
-        metadata = {"id": task_uid, **metadata}
-    if metadata is None:
-        return jsonify({"error": "Job not found"}), 404
-    else:
-        return jsonify(metadata)
-
-
-@app.route("/processed/<task_uid>/<which>")
-def processed(task_uid, which):
-    analysis = get_analysis_data(task_uid)
-    path = analysis.get_processed(which)
-    if not path.exists():
-        return dict(error="Data not found"), 404
-    return send_file(path, mimetype="application/json")
-
-
-@app.route("/logs/<task_uid>/<log_type>")
-def logs(task_uid, log_type):
-    analysis = get_analysis_data(task_uid)
-    path = analysis.get_log(log_type)
-    if not path.exists():
-        return dict(error="Data not found"), 404
-    return send_file(path, mimetype="text/plain")
-
-
-@app.route("/process_info/<task_uid>/<seqid>")
-def process_info(task_uid, seqid):
-    analysis = get_analysis_data(task_uid)
-    if not seqid.isdigit():
-        return dict(error="Incorrect process id"), 400
-    seqid = int(seqid)
-    process_info = analysis.get_process_info(seqid)
-    if process_info is None:
-        return dict(error="Data not found"), 404
-    return jsonify(process_info)
-
-
-@app.route("/logs/<task_uid>/<log_type>/process/<seqid>")
-def process_logs(task_uid, log_type, seqid):
-    analysis = get_analysis_data(task_uid)
-    if not seqid.isdigit():
-        return dict(error="Incorrect process id"), 400
-    seqid = int(seqid)
-    index_path = analysis.get_log_index(f"{log_type}.{seqid}.json")
-    if not index_path.exists():
-        return dict(error="Data not found"), 404
-    index = orjson.loads(index_path.read_text())
-    blocks = index["blocks"]
-    filter_values = request.args.getlist("filter[]")
-    if filter_values:
-        filter_indices = [
-            index["values"].index(filter_value)
-            for filter_value in filter_values
-            if filter_value in index["values"]
-        ]
-        blocks = [
-            block
-            for (block, mapping) in zip(index["blocks"], index["mapping"])
-            if mapping in filter_indices
-        ]
-    log_path = analysis.get_log(log_type)
-    if request.range:
-        if len(request.range.ranges) > 1:
-            return dict(error="Multiple ranges unsupported"), 400
-        range_start, range_stop = request.range.ranges[0]
-        skip = range_start
-        length = (range_stop - range_start + 1) if range_stop is not None else None
-    else:
-        skip, length = 0, None
-    scattered_read = scattered_read_file(log_path, blocks, skip=skip, length=length)
-    return Response(b"".join(scattered_read), mimetype="text/plain")
-
-
-@app.route("/pcap_dump/<task_uid>")
-def pcap_dump(task_uid):
-    """
-    Return archive containing dump.pcap along with extracted tls sessions
-    keys in format acceptable by wireshark.
-    """
-    analysis = get_analysis_data(task_uid)
-    path = analysis.get_pcap_dump()
-    if not path.exists():
-        return dict(error="Data not found"), 404
-    with NamedTemporaryFile() as f_archive:
-        with ZipFile(f_archive, "w", ZIP_DEFLATED) as archive:
-            archive.write(path, "dump.pcap")
-            path = analysis.get_wireshark_key_file()
-            if path.exists():
-                archive.write(path, "dump.keys")
-        f_archive.seek(0)
-        return send_file(f_archive.name, mimetype="application/zip")
-
-
-@app.route("/dumps/<task_uid>")
-def dumps(task_uid):
-    analysis = get_analysis_data(task_uid)
-    path = analysis.get_dumps()
-    if not path.exists():
-        return dict(error="Data not found"), 404
-    return send_file(path, mimetype="application/zip")
-
-
-@app.route("/logs/<task_uid>")
-def list_logs(task_uid):
-    analysis = get_analysis_data(task_uid)
-    return jsonify(list(analysis.list_logs()))
-
-
-@app.route("/graph/<task_uid>")
-def graph(task_uid):
-    analysis = get_analysis_data(task_uid)
-    path = analysis.get_graph()
-    if not path.exists():
-        return dict(error="Data not found"), 404
-    return send_file(path, mimetype="text/plain")
-
-
-@app.route("/screenshot/<task_uid>/<which>")
-def screenshot(task_uid, which):
-    analysis = get_analysis_data(task_uid)
-    if not which.isdigit():
-        return dict(error="Wrong screenshot number"), 400
-    which = int(which)
-    path = analysis.get_screenshot(which)
-    if not path.exists():
-        return dict(error="Data not found"), 404
-    return send_file(path, mimetype="image/png")
-
-
 @app.route("/")
 def index():
     return send_file("frontend/dist/index.html")

drakrun/web/frontend/src/AnalysisScreenshotViewer.jsx

@@ -4,6 +4,8 @@ import "react-medium-image-zoom/dist/styles.css";
 let BASE_URL = "";
 if (import.meta.env.VITE_API_SERVER) {
     BASE_URL = import.meta.env.VITE_API_SERVER;
+} else {
+    BASE_URL = "/api";
 }
 
 export function AnalysisScreenshotViewer({ analysis }) {

drakrun/web/frontend/src/App.jsx

@@ -43,6 +43,12 @@ function AppSidenav() {
                             Upload sample
                         </Link>
                         <div className="sb-sidenav-menu-heading">Sandbox</div>
+                        <a className="nav-link" href="/openapi/swagger">
+                            <div className="sb-nav-link-icon">
+                                <FontAwesomeIcon icon={faGear} />
+                            </div>
+                            API docs
+                        </a>
                         <a className="nav-link" href="/rq">
                             <div className="sb-nav-link-icon">
                                 <FontAwesomeIcon icon={faGear} />

drakrun/web/frontend/src/api.js

@@ -2,6 +2,8 @@ import axios from "axios";
 
 if (import.meta.env.VITE_API_SERVER) {
     axios.defaults.baseURL = import.meta.env.VITE_API_SERVER;
+} else {
+    axios.defaults.baseURL = "/api";
 }
 
 export async function getAnalysisList({ abortController }) {

drakrun/web/schema.py

@@ -0,0 +1,64 @@
+import uuid
+from typing import Annotated, List, Optional
+
+from flask_openapi3 import FileStorage
+from pydantic import AfterValidator, BaseModel, Field, RootModel
+
+
+class APIErrorResponse(BaseModel):
+    error: str = Field(description="Error message")
+
+
+class UploadFileForm(BaseModel):
+    file: FileStorage
+    timeout: Optional[int] = Field(default=None, description="Analysis timeout")
+    file_name: Optional[str] = Field(default=None, description="Target file name")
+    start_command: Optional[str] = Field(default=None, description="Start command")
+    plugins: Optional[str] = Field(
+        default=None, description="Plugins to use (in JSON array string)"
+    )
+
+
+class UploadAnalysisResponse(BaseModel):
+    task_uid: str = Field(description="Unique analysis ID")
+
+
+class AnalysisResponse(BaseModel):
+    id: str = Field(description="Unique analysis ID")
+    status: str = Field(description="Analysis status")
+    time_started: Optional[str] = Field(
+        default=None, description="Analysis start time in ISO format"
+    )
+    time_ended: Optional[str] = Field(
+        default=None, description="Analysis end time in ISO format"
+    )
+
+
+AnalysisListResponse = RootModel[List[AnalysisResponse]]
+
+
+class AnalysisRequestPath(BaseModel):
+    task_uid: Annotated[str, AfterValidator(lambda x: str(uuid.UUID(x, version=4)))] = (
+        Field(description="Unique analysis ID")
+    )
+
+
+class ProcessedRequestPath(AnalysisRequestPath):
+    which: str
+
+
+class LogsRequestPath(AnalysisRequestPath):
+    log_type: str
+
+
+class ProcessInfoRequestPath(AnalysisRequestPath):
+    seqid: int
+
+
+class ProcessLogsRequestPath(AnalysisRequestPath):
+    log_type: str
+    seqid: int
+
+
+class ScreenshotRequestPath(AnalysisRequestPath):
+    which: int

requirements.txt

@@ -10,6 +10,7 @@ tomli==2.2.1
 # Web analyzer dependencies
 rq==2.3.1
 Flask==3.0.3
+flask-openapi3[swagger]==4.1.0
 python-magic==0.4.27
 rq-dashboard==0.8.2.2
 # Screenshotter dependencies