#!/usr/bin/env python3
"""
Usage:
    python3 auto_solve.py --url <...>
"""

import argparse
import json
import os
import re
import subprocess
import sys
import tempfile
import urllib.request
import zipfile
from pathlib import Path, PurePosixPath
from stat import S_IXUSR, S_IXGRP, S_IXOTH

HOOK_SOURCE = r"""
#define _GNU_SOURCE
#include <dlfcn.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

typedef int (*strcmp_t)(const char *, const char *);
typedef int (*memcmp_t)(const void *, const void *, size_t);

static __thread int rc_in_hook = 0;

int strcmp(const char *s1, const char *s2) {
    static strcmp_t real_strcmp = NULL;
    if (!real_strcmp) {
        real_strcmp = (strcmp_t)dlsym(RTLD_NEXT, "strcmp");
    }
    if (!rc_in_hook && s2) {
        rc_in_hook = 1;
        dprintf(2, "[RC_HOOK_STR] secret=%s\n", s2);
        rc_in_hook = 0;
    }
    return real_strcmp(s1, s2);
}

int memcmp(const void *s1, const void *s2, size_t n) {
    static memcmp_t real_memcmp = NULL;
    if (!real_memcmp) {
        real_memcmp = (memcmp_t)dlsym(RTLD_NEXT, "memcmp");
    }
    if (!rc_in_hook && s2 && n <= 1024) {
        rc_in_hook = 1;
        dprintf(2, "[RC_HOOK_MEM] len=%zu data=", n);
        const unsigned char *p = (const unsigned char *)s2;
        for (size_t i = 0; i < n; ++i) {
            dprintf(2, "%02x", p[i]);
        }
        dprintf(2, "\n");
        rc_in_hook = 0;
    }
    return real_memcmp(s1, s2, n);
}
"""

MEM_RE = re.compile(r"\[RC_HOOK_MEM\]\s*len=(\d+)\s+data=([0-9a-fA-F]+)")
STR_RE = re.compile(r"\[RC_HOOK_STR\]\s*secret=(.+)")
DEFAULT_URL = os.environ.get("RC_URL", "URL")


def normalize_member_name(name: str) -> str:
    return PurePosixPath(name).name


def request_archive(base_url: str, target_dir: Path) -> tuple[str, Path, list[str]]:
    url = base_url.rstrip('/') + '/generate'
    req = urllib.request.Request(url, data=b'', method='POST')
    try:
        with urllib.request.urlopen(req) as resp:
            data = resp.read()
            cookies = resp.headers.get_all('Set-Cookie') or []
    except urllib.error.HTTPError as exc:
        raise RuntimeError(f"/generate failed: {exc.read().decode(errors='ignore')}") from exc

    session_id = None
    for cookie in cookies:
        parts = cookie.split(';')
        for part in parts:
            part = part.strip()
            if part.startswith('session_id='):
                session_id = part.split('=', 1)[1]
                break
        if session_id:
            break
    if not session_id:
        raise RuntimeError('session_id cookie was not returned by the server')

    archive_path = target_dir / 'bundle.zip'
    archive_path.write_bytes(data)
    extract_dir = target_dir / 'crackmes'
    extract_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(archive_path, 'r') as zf:
        archive_order = []
        for member in zf.namelist():
            normalized = normalize_member_name(member)
            if normalized and normalized not in archive_order:
                archive_order.append(normalized)
        zf.extractall(extract_dir)

    archive_order = [name for name in archive_order if name in EXTRACTORS]
    if len(archive_order) != len(EXTRACTORS):
        missing = set(EXTRACTORS) - set(archive_order)
        raise RuntimeError(f"unexpected crackme list in archive, missing: {sorted(missing)}")
    
    for item in extract_dir.iterdir():
        if item.is_file():
            mode = item.stat().st_mode | S_IXUSR | S_IXGRP | S_IXOTH
            item.chmod(mode)

    return session_id, extract_dir, archive_order


def build_hook(temp_dir: Path) -> Path:
    src = temp_dir / 'rc_hook.c'
    so = temp_dir / 'rc_hook.so'
    src.write_text(HOOK_SOURCE)
    compile_cmd = ['gcc', '-shared', '-fPIC', '-O2', str(src), '-o', str(so), '-ldl']
    try:
        subprocess.run(compile_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except FileNotFoundError as exc:
        raise RuntimeError('gcc is required to build the hook but was not found') from exc
    return so


def run_binary(binary: Path, hook: Path, payload: bytes, timeout: float = 3.0) -> str:
    env = os.environ.copy()
    env['LD_PRELOAD'] = str(hook)
    proc = subprocess.run(
        [str(binary)],
        input=payload,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env,
        timeout=timeout,
        check=False,
    )
    return proc.stderr.decode(errors='ignore')


def extract_pass1(binary: Path, hook: Path) -> str:
    logs = run_binary(binary, hook, b'test\n')
    m = STR_RE.search(logs)
    if not m:
        raise RuntimeError('Failed to capture strcmp secret for crackme1')
    return m.group(1).strip()


def extract_mem_secret(binary: Path, hook: Path, payload: bytes) -> tuple[int, bytes]:
    logs = run_binary(binary, hook, payload)
    m = MEM_RE.search(logs)
    if not m:
        raise RuntimeError(f'No memcmp hook output for {binary.name}')
    length = int(m.group(1))
    data = bytes.fromhex(m.group(2))
    return length, data[:length]


def extract_pass2(binary: Path, hook: Path) -> str:
    for length in range(1, 33):
        try:
            _, data = extract_mem_secret(binary, hook, ("A" * length + "\n").encode())
            break
        except RuntimeError:
            continue
    else:
        raise RuntimeError('Unable to trigger memcmp for crackme2')

    key = b'cyber32'
    plaintext = bytes(data[i] ^ key[i % len(key)] for i in range(len(data)))
    return plaintext.decode()


def extract_pass3(binary: Path, hook: Path) -> str:
    payload = b"0" * 64 + b"\n"
    logs = run_binary(binary, hook, payload)


    matches = list(MEM_RE.finditer(logs))
    if not matches:
        raise RuntimeError("No memcmp detected in crackme3")

    for match in reversed(matches):
        length = int(match.group(1))
        hex_data = match.group(2)
        if length == 32 and len(hex_data) == 64:
            return hex_data.lower()


    last = matches[-1]
    length = int(last.group(1))
    data = bytes.fromhex(last.group(2))[:length]
    return data.hex()


def extract_pass4(binary: Path, hook: Path) -> str:
    _, data = extract_mem_secret(binary, hook, b'A' * 4)
    return data.hex()


EXTRACTORS = {
    'crackme1': extract_pass1,
    'crackme2': extract_pass2,
    'crackme3': extract_pass3,
    'crackme4': extract_pass4,
}


def solve_in_order(crackme_dir: Path, hook: Path, archive_order: list[str]) -> list[str]:
    answers = []
    for name in archive_order:
        extractor = EXTRACTORS.get(name)
        if not extractor:
            raise RuntimeError(f'unknown crackme entry: {name}')
        binary = crackme_dir / name
        if not binary.exists():
            raise RuntimeError(f'binary {binary} was not extracted')
        answers.append(extractor(binary, hook))
    return answers


def submit_answer(base_url: str, session_id: str, combined: str) -> dict:
    url = base_url.rstrip('/') + '/submit'
    payload = json.dumps({'answer': combined}).encode()
    headers = {
        'Content-Type': 'application/json',
        'Cookie': f'session_id={session_id}',
    }
    req = urllib.request.Request(url, data=payload, method='POST', headers=headers)
    try:
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as exc:
        raise RuntimeError(f'/submit failed: {exc.read().decode(errors="ignore")}') from exc


def main():
    parser = argparse.ArgumentParser(description='Automatic solver for ReverseConveyor')
    parser.add_argument('--url', default=DEFAULT_URL, help='Base URL of the service (default: %(default)s)')
    args = parser.parse_args()

    with tempfile.TemporaryDirectory() as tmp:
        tmp_path = Path(tmp)
        session_id, crackme_dir, archive_order = request_archive(args.url, tmp_path)
        hook = build_hook(tmp_path)

        print(f"[+] Archive order: {' '.join(archive_order)}")
        parts = solve_in_order(crackme_dir, hook, archive_order)

        combined = '_'.join(parts)
        print(f'[+] Answer: {combined}')
        result = submit_answer(args.url, session_id, combined)
        if result.get('ok'):
            print(f"[+] Flag: {result['flag']}")
        else:
            print('[-] Server rejected the answer:', result)


if __name__ == '__main__':
    try:
        main()
    except Exception as exc:
        print(f'[!] {exc}', file=sys.stderr)
        sys.exit(1)