from pathlib import Path
import asyncio
import random
import string
import time

# run cargo build or cargo build --release before this

OK = 101
CORRUPT = 102
MUMBLE = 103
DOWN = 104
CHECKER_ERROR = 110


RELEASE = True
EXPECTED_CHECK_CODE = OK
EXPECTED_PUT_CODE = OK
EXPECTED_GET_CODE = OK
CHECKS_PER_SECOND = 5
CHECKS_PER_REPORT = 100
MAX_PARALLEL_CHECKS = CHECKS_PER_SECOND * 3
VULNS = 1
TIMEOUT = 30
HOST = "127.0.0.1"

assert VULNS == 1, "multiple vulns not supported."

check_period = 1 / CHECKS_PER_SECOND

if RELEASE:
    checker_binary_path = Path("../target/release/checker")
else:
    checker_binary_path = Path("../target/debug/checker")

detailed_log_filepath = Path("detailed.log")

def returncode_to_string(returncode):
    if returncode == 101:
        return "OK"
    elif returncode == 102:
        return "CORRUPT"
    elif returncode == 103:
        return "MUMBLE"
    elif returncode == 104:
        return "DOWN"
    elif returncode == 110:
        return "CHECKER_ERROR"
    else:
        return "UNEXPECTED"

def get_random_flag():
    return "S" + "".join(random.choices(string.ascii_uppercase + string.digits, k=30)) + "="

output_lock = asyncio.locks.Lock()
async def write_to_detailed_log(entry, *, also_print=False):
    async with output_lock:
        with open(detailed_log_filepath, "a") as f:
            f.write(entry + "\n")
    if also_print:
        print(entry)

assert checker_binary_path.exists(), "checker binary doesn't exist. run cargo build or cargo build --release before this."

async def do_all_checks_once():
    flag = get_random_flag()

    check_process = await asyncio.create_subprocess_exec(
        checker_binary_path, "check", HOST,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    try:
        stdout_data, stderr_data = await asyncio.wait_for(check_process.communicate(), timeout=TIMEOUT)
    except asyncio.TimeoutError:
        stdout_data, stderr_data = await check_process.stdout.read(), await check_process.stderr.read()
        await write_to_detailed_log(
            f"CHECK({flag}): Stuck for {TIMEOUT} seconds. Make sure the checker code has it's own timeouts. stdout={stdout_data} stderr={stderr_data}",
            also_print=True
        )
        check_process.kill()
    else:
        returncode = check_process.returncode
        await write_to_detailed_log(
            f"CHECK({flag}): {returncode_to_string(returncode)} ({returncode}) stdout={stdout_data} stderr={stderr_data}",
            also_print=(EXPECTED_CHECK_CODE is not None and returncode != EXPECTED_CHECK_CODE)
        )

    put_process = await asyncio.create_subprocess_exec(
        checker_binary_path, "put", HOST, "", flag, "1",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    try:
        stdout_data, stderr_data = await asyncio.wait_for(put_process.communicate(), timeout=TIMEOUT)
    except asyncio.TimeoutError:
        stdout_data, stderr_data = await put_process.stdout.read(), await put_process.stderr.read()
        await write_to_detailed_log(
            f"PUT({flag}): Stuck for {TIMEOUT} seconds. Make sure the checker code has it's own timeouts. stdout={stdout_data} stderr={stderr_data}",
            also_print=True
        )
        put_process.kill()
        failed = True
    else:
        returncode = put_process.returncode
        await write_to_detailed_log(
            f"PUT({flag}): {returncode_to_string(returncode)} ({returncode}) stdout(public)={stdout_data} stderr(private)={stderr_data}",
            also_print=(EXPECTED_PUT_CODE is not None and returncode != EXPECTED_PUT_CODE)
        )
        flag_id = stderr_data
        failed = returncode != OK

    if failed:
        await write_to_detailed_log(
            f"GET({flag}): SKIPPED since PUT failed.",
            also_print=(EXPECTED_GET_CODE is not None)
        )
        return

    get_process = await asyncio.create_subprocess_exec(
        checker_binary_path, "get", HOST, flag_id, flag, "1",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    try:
        stdout_data, stderr_data = await asyncio.wait_for(get_process.communicate(), timeout=TIMEOUT)
    except asyncio.TimeoutError:
        stdout_data, stderr_data = await get_process.stdout.read(), await get_process.stderr.read()
        await write_to_detailed_log(
            f"GET({flag}): Stuck for {TIMEOUT} seconds. Make sure the checker code has it's own timeouts. stdout={stdout_data} stderr={stderr_data}",
            also_print=True
        )
        get_process.kill()
    else:
        returncode = get_process.returncode
        await write_to_detailed_log(
            f"GET({flag}): {returncode_to_string(returncode)} ({returncode}) stdout(public)={stdout_data} stderr(private)={stderr_data}",
            also_print=(EXPECTED_GET_CODE is not None and returncode != EXPECTED_GET_CODE)
        )


async def main():
    pending_tasks = []
    checks_done = 0
    while True:
        if len(pending_tasks) > 0 and pending_tasks[0].done() and (checks_done + 1) % CHECKS_PER_REPORT == 0:
            print(f"Already done {checks_done + 1} checks. {len(pending_tasks) - 1} running in parallel.")
        while len(pending_tasks) > 0 and pending_tasks[0].done():
            pending_tasks.pop(0)
            checks_done += 1
        if len(pending_tasks) <= MAX_PARALLEL_CHECKS:
            pending_tasks.append(asyncio.create_task(asyncio.wait_for(do_all_checks_once(), timeout=TIMEOUT * 3 + 10)))
        else:
            print(f"Throttling checks since there are at least {MAX_PARALLEL_CHECKS} running")
        await asyncio.sleep(check_period)


if __name__ == "__main__":
    asyncio.run(main())