Before / After Code Comparison

Before remediation

1from __future__ import annotations23import json4import os5import queue6import sqlite37import threading8from datetime import datetime9from pathlib import Path1011from flask import Flask, Response, g, jsonify, request, send_file1213BASE_DIR = Path(__file__).resolve().parent14DATA_DIR = BASE_DIR / "data"15STORAGE_DIR = DATA_DIR / "storage"16DB_PATH = DATA_DIR / "docs.db"1718DATA_DIR.mkdir(parents=True, exist_ok=True)19STORAGE_DIR.mkdir(parents=True, exist_ok=True)2021app = Flask(__name__)22job_queue: "queue.Queue[dict]" = queue.Queue()232425def get_db() -> sqlite3.Connection:26    if "db" not in g:27        conn = sqlite3.connect(DB_PATH)28        conn.row_factory = sqlite3.Row29        g.db = conn30    return g.db313233def init_db() -> None:34    db = sqlite3.connect(DB_PATH)35    db.executescript(36        """37        create table if not exists organizations (38            id integer primary key autoincrement,39            name text not null,40            org_code text not null unique,41            region text not null,42            tier text not null43        );44        create table if not exists teams (45            id integer primary key autoincrement,46            org_id integer not null,47            name text not null,48            cost_center text not null,49            status text not null default 'active'50        );51        create table if not exists users (52            id integer primary key autoincrement,53            org_id integer not null,54            team_id integer,55            email text not null,56            display_name text not null,57            role text not null,58            status text not null default 'active'59        );60        create table if not exists accounts (61            id integer primary key autoincrement,62            org_id integer not null,63            account_no text not null,64            currency text not null,65            balance numeric not null default 066        );67        create table if not exists beneficiaries (68            id integer primary key autoincrement,69            org_id integer not null,70            name text not null,71            bank_name text not null,72            routing text not null73        );74        create table if not exists transfers (75            id integer primary key autoincrement,76            org_id integer not null,77            account_id integer not null,78            beneficiary_id integer not null,79            amount numeric not null,80            currency text not null,81            status text not null,82            memo text not null,83            requested_by text not null84        );85        create table if not exists documents (86            id integer primary key autoincrement,87            org_id integer not null,88            filename text not null,89            original_name text not null,90            mime_type text not null,91            kind text not null,92            created_at text not null default current_timestamp93        );94        create table if not exists risk_reviews (95            id integer primary key autoincrement,96            org_id integer not null,97            subject text not null,98            score numeric not null,99            note text not null,100            created_at text not null default current_timestamp101        );102        create table if not exists audit_records (103            id integer primary key autoincrement,104            org_id integer,105            actor text not null,106            action text not null,107            resource text not null,108            details text not null,109            created_at text not null default current_timestamp110        );111        """112    )113    if db.execute("select count(*) from organizations").fetchone()[0] == 0:114        db.execute(115            "insert into organizations (name, org_code, region, tier) values (?, ?, ?, ?)",116            ("Acme Treasury", "acme", "us-east", "enterprise"),117        )118        db.execute(119            "insert into organizations (name, org_code, region, tier) values (?, ?, ?, ?)",120            ("Northwind Holdings", "northwind", "eu-west", "growth"),121        )122        db.commit()123    db.close()124125126def current_identity() -> dict:127    return {128        "org_id": int(request.headers.get("X-Org-Id", "1")),129        "user_id": request.headers.get("X-User-Id", "docs-user"),130        "role": request.headers.get("X-Role", "viewer"),131    }132133134@app.before_request135def load_identity() -> None:136    g.identity = current_identity()137    g.started_at = datetime.utcnow().isoformat()138139140@app.after_request141def audit_request(response: Response) -> Response:142    identity = getattr(g, "identity", None)143    if identity:144        db = get_db()145        db.execute(146            "insert into audit_records (org_id, actor, action, resource, details) values (?, ?, ?, ?, ?)",147            (148                identity["org_id"],149                identity["user_id"],150                f"{request.method} {request.path}",151                request.full_path,152                json.dumps({"status": response.status_code, "args": request.args.to_dict(), "json": request.get_json(silent=True)}),153            ),154        )155        db.commit()156    return response157158159@app.teardown_appcontext160def close_db(exc: Exception | None) -> None:161    db = g.pop("db", None)162    if db is not None:163        db.close()164165166def require_roles(*allowed: str):167    def decorator(fn):168        def wrapped(*args, **kwargs):169            if g.identity["role"] not in allowed:170                return jsonify({"error": "forbidden"}), 403171            return fn(*args, **kwargs)172173        wrapped.__name__ = fn.__name__174        return wrapped175176    return decorator177178179def render_report(org_name: str, title: str, notes: str) -> str:180    return f"""181    <html>182      <body>183        <h1>{title}</h1>184        <h2>{org_name}</h2>185        <div>{notes}</div>186      </body>187    </html>188    """189190191def enqueue_job(payload: dict) -> None:192    job_queue.put(payload)193194195def job_worker() -> None:196    while True:197        payload = job_queue.get()198        if payload is None:199            break200        db = sqlite3.connect(DB_PATH)201        db.execute(202            "insert into audit_records (org_id, actor, action, resource, details) values (?, ?, ?, ?, ?)",203            (204                payload.get("org_id"),205                "system",206                "job.dispatch",207                payload.get("kind", "document"),208                json.dumps(payload),209            ),210        )211        if payload.get("kind") == "document":212            filename = payload["filename"]213            target = STORAGE_DIR / filename214            target.write_text(payload.get("content", ""), encoding="utf-8")215            db.execute(216                "insert into documents (org_id, filename, original_name, mime_type, kind) values (?, ?, ?, ?, ?)",217                (218                    payload.get("org_id"),219                    filename,220                    payload.get("original_name", filename),221                    payload.get("mime_type", "text/plain"),222                    payload.get("kind", "document"),223                ),224            )225        db.commit()226        db.close()227        job_queue.task_done()228229230init_db()231threading.Thread(target=job_worker, daemon=True).start()232233234@app.route("/health", methods=["GET"])235def health() -> tuple:236    db = get_db()237    counts = {238        "documents": db.execute("select count(*) from documents").fetchone()[0],239        "audit_records": db.execute("select count(*) from audit_records").fetchone()[0],240        "jobs_pending": job_queue.qsize(),241    }242    return jsonify({"ok": True, "counts": counts})243244245@app.route("/documents", methods=["GET"])246@require_roles("admin", "manager", "viewer")247def search_documents() -> tuple:248    clauses = ["org_id = ?"]249    params: list[object] = [g.identity["org_id"]]250    kind = request.args.get("kind", "").strip()251    name = request.args.get("name", "").strip()252    if kind:253        clauses.append("kind = ?")254        params.append(kind)255    if name:256        clauses.append("original_name like ?")257        params.append(f"%{name}%")258    rows = get_db().execute(259        f"select id, filename, original_name, mime_type, kind, created_at from documents where {' and '.join(clauses)} order by id desc",260        params,261    ).fetchall()262    return jsonify({"data": [dict(row) for row in rows]})263264265@app.route("/docs/render", methods=["POST"])266@require_roles("admin", "manager")267def render_document() -> tuple:268    payload = request.get_json(force=True)269    org_name = payload["orgName"]270    title = payload.get("title", "Treasury Notice")271    notes = payload.get("notes", "")272    filename = payload.get("filename", f"report-{int(datetime.utcnow().timestamp())}.html")273    html = render_report(org_name, title, notes)274    target = STORAGE_DIR / filename275    target.write_text(html, encoding="utf-8")276    db = get_db()277    db.execute(278        "insert into documents (org_id, filename, original_name, mime_type, kind) values (?, ?, ?, ?, ?)",279        (g.identity["org_id"], filename, filename, "text/html", "report"),280    )281    db.commit()282    return jsonify({"filename": filename, "path": str(target)}), 201283284285@app.route("/docs/upload", methods=["POST"])286@require_roles("admin", "manager")287def upload_document() -> tuple:288    upload = request.files.get("file")289    if upload is None:290        return jsonify({"error": "missing file"}), 400291    filename = f"{int(datetime.utcnow().timestamp())}-{upload.filename}"292    target = STORAGE_DIR / filename293    upload.save(target)294    db = get_db()295    db.execute(296        "insert into documents (org_id, filename, original_name, mime_type, kind) values (?, ?, ?, ?, ?)",297        (g.identity["org_id"], filename, upload.filename, upload.mimetype or "application/octet-stream", "upload"),298    )299    db.commit()300    return jsonify({"filename": filename}), 201301302303@app.route("/docs/download/<path:filename>", methods=["GET"])304@require_roles("admin", "manager", "viewer")305def download_document(filename: str):306    target = STORAGE_DIR / filename307    if not target.exists():308        return jsonify({"error": "not found"}), 404309    return send_file(target, as_attachment=True)310311312@app.route("/jobs/dispatch", methods=["POST"])313@require_roles("admin", "manager")314def dispatch_job() -> tuple:315    payload = request.get_json(force=True)316    payload["org_id"] = g.identity["org_id"]317    enqueue_job(payload)318    return jsonify({"queued": True, "kind": payload.get("kind", "document")}), 202319320321@app.route("/reports/summary", methods=["GET"])322@require_roles("admin", "manager")323def summary_report() -> tuple:324    db = get_db()325    org_id = int(request.args.get("orgId", g.identity["org_id"]))326    title = request.args.get("title", "Document Summary")327    org = db.execute("select name, region, tier from organizations where id = ?", (org_id,)).fetchone()328    if org is None:329        return jsonify({"error": "missing org"}), 404330    notes = request.args.get("notes", "No notes supplied")331    html = render_report(org["name"], title, notes)332    return Response(html, mimetype="text/html")333334335if __name__ == "__main__":336    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", "4300")), debug=False)

After remediation

1from __future__ import annotations23import json4import os5import queue6import sqlite37import threading8from datetime import datetime9from pathlib import Path1011from flask import Flask, Response, g, jsonify, request, send_file1213BASE_DIR = Path(__file__).resolve().parent14DATA_DIR = Path(os.environ.get("DOCS_WORKER_DATA_DIR", str(BASE_DIR / "data"))).resolve()15STORAGE_DIR = DATA_DIR / "storage"16DB_PATH = DATA_DIR / "docs.db"1718DATA_DIR.mkdir(parents=True, exist_ok=True)19STORAGE_DIR.mkdir(parents=True, exist_ok=True)2021app = Flask(__name__)22job_queue: "queue.Queue[dict]" = queue.Queue()232425def get_db() -> sqlite3.Connection:26    if "db" not in g:27        conn = sqlite3.connect(DB_PATH)28        conn.row_factory = sqlite3.Row29        g.db = conn30    return g.db313233def init_db() -> None:34    db = sqlite3.connect(DB_PATH)35    db.executescript(36        """37        create table if not exists organizations (38            id integer primary key autoincrement,39            name text not null,40            org_code text not null unique,41            region text not null,42            tier text not null43        );44        create table if not exists teams (45            id integer primary key autoincrement,46            org_id integer not null,47            name text not null,48            cost_center text not null,49            status text not null default 'active'50        );51        create table if not exists users (52            id integer primary key autoincrement,53            org_id integer not null,54            team_id integer,55            email text not null,56            display_name text not null,57            role text not null,58            status text not null default 'active'59        );60        create table if not exists accounts (61            id integer primary key autoincrement,62            org_id integer not null,63            account_no text not null,64            currency text not null,65            balance numeric not null default 066        );67        create table if not exists beneficiaries (68            id integer primary key autoincrement,69            org_id integer not null,70            name text not null,71            bank_name text not null,72            routing text not null73        );74        create table if not exists transfers (75            id integer primary key autoincrement,76            org_id integer not null,77            account_id integer not null,78            beneficiary_id integer not null,79            amount numeric not null,80            currency text not null,81            status text not null,82            memo text not null,83            requested_by text not null84        );85        create table if not exists documents (86            id integer primary key autoincrement,87            org_id integer not null,88            filename text not null,89            original_name text not null,90            mime_type text not null,91            kind text not null,92            created_at text not null default current_timestamp93        );94        create table if not exists risk_reviews (95            id integer primary key autoincrement,96            org_id integer not null,97            subject text not null,98            score numeric not null,99            note text not null,100            created_at text not null default current_timestamp101        );102        create table if not exists audit_records (103            id integer primary key autoincrement,104            org_id integer,105            actor text not null,106            action text not null,107            resource text not null,108            details text not null,109            created_at text not null default current_timestamp110        );111        """112    )113    if db.execute("select count(*) from organizations").fetchone()[0] == 0:114        db.execute(115            "insert into organizations (name, org_code, region, tier) values (?, ?, ?, ?)",116            ("Acme Treasury", "acme", "us-east", "enterprise"),117        )118        db.execute(119            "insert into organizations (name, org_code, region, tier) values (?, ?, ?, ?)",120            ("Northwind Holdings", "northwind", "eu-west", "growth"),121        )122        db.commit()123    db.close()124125126def current_identity() -> dict:127    return {128        "org_id": int(request.headers.get("X-Org-Id", "1")),129        "user_id": request.headers.get("X-User-Id", "docs-user"),130        "role": request.headers.get("X-Role", "viewer"),131    }132133134@app.before_request135def load_identity() -> None:136    g.identity = current_identity()137    g.started_at = datetime.utcnow().isoformat()138139140@app.after_request141def audit_request(response: Response) -> Response:142    identity = getattr(g, "identity", None)143    if identity:144        db = get_db()145        db.execute(146            "insert into audit_records (org_id, actor, action, resource, details) values (?, ?, ?, ?, ?)",147            (148                identity["org_id"],149                identity["user_id"],150                f"{request.method} {request.path}",151                request.full_path,152                json.dumps({"status": response.status_code, "args": request.args.to_dict(), "json": request.get_json(silent=True)}),153            ),154        )155        db.commit()156    return response157158159@app.teardown_appcontext160def close_db(exc: Exception | None) -> None:161    db = g.pop("db", None)162    if db is not None:163        db.close()164165166def require_roles(*allowed: str):167    def decorator(fn):168        def wrapped(*args, **kwargs):169            if g.identity["role"] not in allowed:170                return jsonify({"error": "forbidden"}), 403171            return fn(*args, **kwargs)172173        wrapped.__name__ = fn.__name__174        return wrapped175176    return decorator177178179def render_report(org_name: str, title: str, notes: str) -> str:180    return f"""181    <html>182      <body>183        <h1>{title}</h1>184        <h2>{org_name}</h2>185        <div>{notes}</div>186      </body>187    </html>188    """189190191def resolve_storage_path(filename: str, *, must_exist: bool = False) -> Path:192    if not isinstance(filename, str) or not filename:193        raise ValueError("invalid filename")194    if filename in {".", ".."}:195        raise ValueError("invalid filename")196    if "\x00" in filename:197        raise ValueError("invalid filename")198    candidate_name = Path(filename).name199    if candidate_name != filename or Path(filename).is_absolute():200        raise ValueError("invalid filename")201202    base_dir = STORAGE_DIR.resolve()203    target = base_dir / filename204205    if must_exist:206        resolved = target.resolve(strict=True)207        if resolved.parent != base_dir:208            raise ValueError("invalid filename")209        return resolved210211    if target.exists():212        resolved = target.resolve(strict=True)213        if resolved.parent != base_dir:214            raise ValueError("invalid filename")215216    return target217218219def enqueue_job(payload: dict) -> None:220    job_queue.put(payload)221222223def job_worker() -> None:224    while True:225        payload = job_queue.get()226        if payload is None:227            break228        db = sqlite3.connect(DB_PATH)229        db.execute(230            "insert into audit_records (org_id, actor, action, resource, details) values (?, ?, ?, ?, ?)",231            (232                payload.get("org_id"),233                "system",234                "job.dispatch",235                payload.get("kind", "document"),236                json.dumps(payload),237            ),238        )239        if payload.get("kind") == "document":240            filename = payload["filename"]241            target = STORAGE_DIR / filename242            target.write_text(payload.get("content", ""), encoding="utf-8")243            db.execute(244                "insert into documents (org_id, filename, original_name, mime_type, kind) values (?, ?, ?, ?, ?)",245                (246                    payload.get("org_id"),247                    filename,248                    payload.get("original_name", filename),249                    payload.get("mime_type", "text/plain"),250                    payload.get("kind", "document"),251                ),252            )253        db.commit()254        db.close()255        job_queue.task_done()256257258init_db()259threading.Thread(target=job_worker, daemon=True).start()260261262@app.route("/health", methods=["GET"])263def health() -> tuple:264    db = get_db()265    counts = {266        "documents": db.execute("select count(*) from documents").fetchone()[0],267        "audit_records": db.execute("select count(*) from audit_records").fetchone()[0],268        "jobs_pending": job_queue.qsize(),269    }270    return jsonify({"ok": True, "counts": counts})271272273@app.route("/documents", methods=["GET"])274@require_roles("admin", "manager", "viewer")275def search_documents() -> tuple:276    clauses = ["org_id = ?"]277    params: list[object] = [g.identity["org_id"]]278    kind = request.args.get("kind", "").strip()279    name = request.args.get("name", "").strip()280    if kind:281        clauses.append("kind = ?")282        params.append(kind)283    if name:284        clauses.append("original_name like ?")285        params.append(f"%{name}%")286    rows = get_db().execute(287        f"select id, filename, original_name, mime_type, kind, created_at from documents where {' and '.join(clauses)} order by id desc",288        params,289    ).fetchall()290    return jsonify({"data": [dict(row) for row in rows]})291292293@app.route("/docs/render", methods=["POST"])294@require_roles("admin", "manager")295def render_document() -> tuple:296    payload = request.get_json(force=True)297    org_name = payload["orgName"]298    title = payload.get("title", "Treasury Notice")299    notes = payload.get("notes", "")300    filename = payload.get("filename", f"report-{int(datetime.utcnow().timestamp())}.html")301    html = render_report(org_name, title, notes)302    try:303        target = resolve_storage_path(filename)304    except ValueError:305        return jsonify({"error": "invalid filename"}), 400306    target.write_text(html, encoding="utf-8")307    db = get_db()308    db.execute(309        "insert into documents (org_id, filename, original_name, mime_type, kind) values (?, ?, ?, ?, ?)",310        (g.identity["org_id"], filename, filename, "text/html", "report"),311    )312    db.commit()313    return jsonify({"filename": filename, "path": str(target)}), 201314315316@app.route("/docs/upload", methods=["POST"])317@require_roles("admin", "manager")318def upload_document() -> tuple:319    upload = request.files.get("file")320    if upload is None:321        return jsonify({"error": "missing file"}), 400322    filename = f"{int(datetime.utcnow().timestamp())}-{upload.filename}"323    try:324        target = resolve_storage_path(filename)325    except ValueError:326        return jsonify({"error": "invalid filename"}), 400327    upload.save(target)328    db = get_db()329    db.execute(330        "insert into documents (org_id, filename, original_name, mime_type, kind) values (?, ?, ?, ?, ?)",331        (g.identity["org_id"], filename, upload.filename, upload.mimetype or "application/octet-stream", "upload"),332    )333    db.commit()334    return jsonify({"filename": filename}), 201335336337@app.route("/docs/download/<path:filename>", methods=["GET"])338@require_roles("admin", "manager", "viewer")339def download_document(filename: str):340    try:341        target = resolve_storage_path(filename, must_exist=True)342    except FileNotFoundError:343        return jsonify({"error": "not found"}), 404344    except ValueError:345        return jsonify({"error": "invalid filename"}), 400346    return send_file(target, as_attachment=True)347348349@app.route("/jobs/dispatch", methods=["POST"])350@require_roles("admin", "manager")351def dispatch_job() -> tuple:352    payload = request.get_json(force=True)353    payload["org_id"] = g.identity["org_id"]354    enqueue_job(payload)355    return jsonify({"queued": True, "kind": payload.get("kind", "document")}), 202356357358@app.route("/reports/summary", methods=["GET"])359@require_roles("admin", "manager")360def summary_report() -> tuple:361    db = get_db()362    org_id = int(request.args.get("orgId", g.identity["org_id"]))363    title = request.args.get("title", "Document Summary")364    org = db.execute("select name, region, tier from organizations where id = ?", (org_id,)).fetchone()365    if org is None:366        return jsonify({"error": "missing org"}), 404367    notes = request.args.get("notes", "No notes supplied")368    html = render_report(org["name"], title, notes)369    return Response(html, mimetype="text/html")370371372if __name__ == "__main__":373    app.run(host="127.0.0.1", port=int(os.environ.get("PORT", "4300")), debug=False)

Before remediation

1Flask==3.0.3

After remediation

1Flask==3.1.3

Before remediation

missing

After remediation

1import os2import shutil3import tempfile4import importlib5import runpy6import unittest7from pathlib import Path8from urllib.parse import quote9from unittest.mock import patch101112TEST_DATA_DIR = tempfile.mkdtemp(prefix="docs-worker-test-")13os.environ["DOCS_WORKER_DATA_DIR"] = TEST_DATA_DIR1415docs_app = importlib.import_module("app")161718class DownloadDocumentTests(unittest.TestCase):19    @classmethod20    def setUpClass(cls):21        docs_app.app.config["TESTING"] = True22        cls.client = docs_app.app.test_client()23        cls.storage_dir = Path(TEST_DATA_DIR) / "storage"2425    @classmethod26    def tearDownClass(cls):27        shutil.rmtree(TEST_DATA_DIR, ignore_errors=True)2829    def test_download_rejects_traversal_and_symlink_escape(self):30        outside = Path(TEST_DATA_DIR) / "outside.txt"31        outside.write_text("secret", encoding="utf-8")32        self.storage_dir.mkdir(parents=True, exist_ok=True)33        (self.storage_dir / "escape.txt").symlink_to(outside)3435        traversal = self.client.get(36            f"/docs/download/{quote('../outside.txt', safe='')}",37            headers={"X-Role": "viewer", "X-Org-Id": "1"},38        )39        self.assertEqual(traversal.status_code, 400)40        self.assertEqual(traversal.get_json(), {"error": "invalid filename"})4142        symlink = self.client.get(43            "/docs/download/escape.txt",44            headers={"X-Role": "viewer", "X-Org-Id": "1"},45        )46        self.assertEqual(symlink.status_code, 400)47        self.assertEqual(symlink.get_json(), {"error": "invalid filename"})4849    def test_download_returns_existing_file(self):50        self.storage_dir.mkdir(parents=True, exist_ok=True)51        target = self.storage_dir / "statement.txt"52        target.write_text("document-body", encoding="utf-8")5354        response = self.client.get(55            "/docs/download/statement.txt",56            headers={"X-Role": "viewer", "X-Org-Id": "1"},57        )58        self.assertEqual(response.status_code, 200)59        self.assertEqual(response.data, b"document-body")6061    def test_main_binds_to_loopback(self):62        with patch("flask.Flask.run") as run_mock:63            runpy.run_module("app", run_name="__main__")6465        self.assertTrue(run_mock.called)66        _, kwargs = run_mock.call_args67        self.assertEqual(kwargs["host"], "127.0.0.1")68        self.assertEqual(kwargs["debug"], False)697071if __name__ == "__main__":72    unittest.main()

Before remediation

missing