Speedrun Challenge from Team Oceania Qualifiers 2023

Dockerfile

FROM rust

RUN apt update && apt install -y python3-pip

RUN pip3 install flask

RUN rustup target add wasm32-unknown-unknown

WORKDIR /srv

COPY . .

RUN mkdir /srv/tmp

ENV FLAG oiccflag{D0N0774UN7H4PPYFUN6r4PH63N3r470r4ND5P33DrUNN3r}

ENV FLASK_APP /srv/server.py

CMD ["flask", "run", "-h", "0.0.0.0"]

main.py

from subprocess import check_output
import time
from flask import Flask, render_template, request
import secrets
import os
import graph
import random
from os.path import join

TMP_DIR = "tmp"


def write_file(filename, content):
    with open(filename, "w") as f:
        f.write(content)


def read_file(filename):
    with open(filename, "rb") as f:
        return f.read()


app = Flask(__name__)

# A dictionary to store the tokens
tokens = {}


# Generate a new token for the user
def generate_token():
    return secrets.token_hex(16)


# Generate a new secret code for the user
def generate_code():
    return secrets.token_hex(15)


@app.route("/")
def home():
    # Generate a new token for the user
    token = generate_token()

    end_time = time.time() + 30

    # Save the token in the dictionary
    tokens[token] = (generate_code(), end_time)

    # Include the token in the response
    return render_template("home.html", time=end_time, token=token)


def choice(a): return random.choice(a)


def make_other(a):
    """
    Return a random int between 0 and 255 that is not a
    """

    while True:
        b = random.randint(0, 255)
        if b != a:
            return b


def generate_program(code):
    for i in range(0, 100):
        try:
            return generate_program_impl(code)
        except:
            pass


def generate_program_impl(code):
    g = graph.generate()

    code_i = int(code, base=16)

    # print(hex(code_i))

    funcs = []

    visited = set()

    stubs = set()

    def generate(code: int, node: str, depth: int):
        current_code = code & 0xff
        next_code = code >> 8

        # print(hex(current_code), hex(next_code))

        if node in visited:
            raise NotImplementedError()

        visited.add(node)

        if depth == 15:
            funcs.append(f"""fn {node}(_v: u128) -> u8 {{ 0x1 }}""")
            return

        edges = g.get_edges(node)

        correct = choice(edges)

        func_body = ""

        for edge in edges:
            if edge == correct:
                func_body += f"if c == {hex(current_code)} {{ return {edge}(n) }} // correct\n"
                generate(next_code, edge, depth + 1)
            else:
                other = make_other(current_code)
                func_body += f"if c == {hex(other)} {{ return {edge}(n) }}\n"
                stubs.add(edge)

        funcs.append(f"""fn {node}(v: u128) -> u8 {{
    let c = v & 0xff;
    let n = v >> 8;
    {func_body}
    return 0x0
}}""")

    generate(code_i, "n0", 0)

    i = 1

    def generate_stub(node: str, i: int):
        if node in visited:
            return
        visited.add(node)

        funcs.append(f"""fn {node}(_v: u128) -> u8 {{ {hex(i)} }}""")

    for stub in stubs:
        i += 1
        generate_stub(stub, i)

    return """#![no_std]

use core::panic::PanicInfo;

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    loop {}
}


#[no_mangle]
pub unsafe extern "C" fn exec(input: u128) -> u8 {
    n0(input)
}""" + "\n\n".join(funcs)


def create_program(code: str) -> bytes:
    tmp_name = "tmp_" + secrets.token_hex(16)

    program_code = generate_program(code)

    tmp_name = join(TMP_DIR, tmp_name)

    write_file(tmp_name + ".rs", program_code)

    check_output(["rustc",
                  "-O",
                  tmp_name + ".rs",
                  "-o", tmp_name + ".wasm",
                  "--target", "wasm32-unknown-unknown",
                  "--edition", "2018",
                  "--crate-type", "cdylib"])

    content = read_file(tmp_name + ".wasm")

    os.remove(tmp_name + ".rs")
    os.remove(tmp_name + ".wasm")

    return content


@app.route("/download/<token>", methods=["GET"])
def download_program(token):
    if token in tokens:
        code, timer = tokens[token]
        return create_program(code)
    else:
        return "Invalid Token"


@app.route("/stop_timer", methods=["POST"])
def stop_timer():
    # Get the token from the form
    token = request.form["token"]

    # Get the code from the form
    submitted_code = request.form["code"]

    # Check if the token is valid
    if token in tokens:
        # Check if the timer has not expired
        code, timer = tokens[token]

        if time.time() < timer:
            # Stop the timer
            del tokens[token]

            if submitted_code == code:
                return os.environ.get("FLAG", "placeholder")
            else:
                return "Incorrect Code. Try Again."
        else:
            # Stop the timer
            del tokens[token]

            return "Timer has expired"
    else:
        return "Invalid token"


if __name__ == '__main__':
    print(generate_program("18a623d8b49502a67cba41978698b6"))

templates_home.html

<!DOCTYPE html>
<html>
  <head>
    <title>Speedrun</title>
    <style>
      /* Add your CSS styles here */
      /* Add general styles for the page */
      body {
        background-color: #f5f5f5;
        font-family: Arial, sans-serif;
        margin: 0;
      }
      /* Add styles for the form */
      .form-container {
        background-color: white;
        border: 1px solid #ccc;
        border-radius: 10px;
        box-shadow: 2px 2px 5px rgba(0, 0, 0, 0.3);
        margin: 100px auto;
        max-width: 500px;
        padding: 20px;
        text-align: center;
      }
      /* Add styles for the form inputs */
      .form-container input[type="text"] {
        border: 1px solid #ccc;
        border-radius: 5px;
        padding: 8px;
        width: 80%;
      }
      /* Add styles for the form label */
      .form-container label {
        font-size: 1.2em;
        font-weight: bold;
        margin-bottom: 10px;
      }
      /* Add styles for the submit button */
      .btn {
        display: inline-block;
        background-color: #b8b8b8;
        color: white;
        font-weight: bold;
        padding: 2px;
        border: 1px solid black;
      }

      .btn-primary {
        background-color: #4caf50;
        padding: 10px;
        font-size: 1.2em;
      }
    </style>
  </head>
  <body>
    <div class="form-container">
      <h1>Speedrun</h1>

      <p>You have until <span id="end"></span> to solve this challenge.</p>

      <p>
        All you need to do is download
        <a class="btn" href="/download/{{token}}" target="_blank"
          >this program</a
        >
        and find the secret flag.
      </p>

      <form action="/stop_timer" method="post">
        <input type="hidden" name="token" value="{{token}}" />
        <label for="code">Enter code to stop timer:</label>
        <br />
        <input type="text" id="code" name="code" placeholder="Enter code" />
        <br /><br />
        <input class="btn btn-primary" type="submit" value="Stop Timer" />
      </form>
    </div>

    <script>
      const endTime = new Date(Number.parseFloat("{{time}}") * 1000);

      document.querySelector("#end").textContent = endTime.toString();
    </script>
  </body>
</html>