#!/usr/bin/env python3
from enochecker import BaseChecker, BrokenServiceException, EnoException, run
from enochecker.utils import SimpleSocket, assert_equals, assert_in
import random, string, struct, logging
import numpy as np

logging.getLogger("faker").setLevel(logging.WARNING)
from faker import Faker

fake = Faker(["en_US"])

class STLDoctorChecker(BaseChecker):
    flag_variants = 1
    noise_variants = 1
    havoc_variants = 4
    service_name = "stldoctor"
    port = 9000

    def login_user(self, conn: SimpleSocket, password: str):
        self.debug("Sending command to login.")
        conn.write(f"login\n{password}\n")
        conn.readline_expect(b"logged in!", read_until=b"$", exception_message="Failed to log in")

    def binwrite(conn: SimpleSocket, buf: bytes):
        conn.sock.sendall(buf)

    def fake_filename(self):
        allowed = "abcdefghijklmopqrstuvwxyz0123456789-+.!"
        return "".join([c for c in fake.name().lower().replace(" ", "-") if c in allowed][:60]).ljust(5, "!")

    def generate_ascii_file(self, solidname: str):
        if solidname != "":
            content = f"solid {solidname}\n"
        else:
            content = "solid\n"
        facet_count = int(random.random() * 30) + 4
        for fi in range(facet_count):
            content += "facet normal "
            vs = [[random.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]) + "\n"
            content += "outer loop\n"
            for i in range(3):
                content += "vertex " + " ".join([f"{v:.2f}" for v in vs[i]]) + "\n"
            content += "endloop\n"
            content += "endfacet\n"
        if solidname != b"":
            content += f"endsolid {solidname}\n"
        else:
            content += "endsolid\n"

        return content.encode("latin1")

    def generate_bin_file(self, solidname: str):
        if len(solidname.encode()) > 78:
            raise EnoException("Solidname to embed in header is larger than header itself")
        if solidname != "":
            content = b"#" + solidname.encode("ascii").ljust(78, b"\x00") + b"\x00"
        else:
            content = b"#" + b"\x00" * 79
        facet_count = int(random.random() * 30) + 4
        content += struct.pack("<I", facet_count)
        for fi in range(facet_count):
            vs = [[random.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]))
            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"
        return content

    def generate_file(self, filetype: str, solidname: str):
        if filetype == "ascii":
            return self.generate_ascii_file(solidname = solidname)
        elif filetype == "bin":
            return self.generate_bin_file(solidname = solidname)
        else:
            raise EnoException("Invalid file type supplied");

    def putfile(self, conn: SimpleSocket, solidname: str, modelname: str, filetype: str):
        # Generate file contents
        stlfile = self.generate_file(filetype = filetype, solidname = solidname)

        # Upload file
        self.debug("Sending command to submit file")
        conn.write("submit\n")
        conn.write(f"{len(stlfile)}\n")
        binwrite(conn, stlfile)
        conn.write(f"{modelname}\n")
        self.debug(b":::RESPONSE:::\n" + conn.read_until(b"with ID "))

        # Parse ID
        fileid = conn.read_until(b"!")
        if fileid == b"":
            raise BrokenServiceException("Unable to upload file!")
        self.debug(f"Got ID {fileid}")

        conn.read_until(b"$")

        return stlfile, fileid

    def getfile(self, conn: SimpleSocket, modelname: str):
        if modelname != "":
            self.debug(f"Sending command to retrieve file with '{modelname}'")
            conn.write(f"query\n{modelname}\n0\ny\n")
        else:
            self.debug(f"Sending command to retrieve file")
            conn.write(f"query\n0\ny\n")

        resp = conn.read_until(b"$")

        return resp

    def querydb(self, *args):
        self.debug("Querying db contents");
        vals = []
        for arg in args:
            try:
                val: str = self.chain_db[arg]
            except IndexError as ex:
                raise BrokenServiceException("Invalid db contents")
            vals.append(val)
        return vals

    def postdb(self, vdict):
        self.chain_db = vdict

    def havoc_upload(self, filetype: str, registered: bool):
        conn = self.openconn()
        if registered:
            pass # TODO: auth
        modelname = self.fake_filename()[:50]
        modelname += "".join([chr(127 + int(random.random() * 128)) for i in range(10)]) # noise
        contents, mid = self.putfile(conn, filetype = filetype, modelname = modelname, solidname = self.fake_filename())
        resp = self.getfile(conn, modelname = modelname)
        assert_in(modelname.encode(), resp, f"Model name '{modelname}' not returned / correctly parsed")
        assert_in(solidname.encode(), resp, f"Solid name '{modelname}' not returned / correctly parsed")
        assert_in(contents, resp, f"STL File contents not returned / correctly parsed")
        self.closeconn(conn)

        conn = self.openconn()
        resp = self.getfile(conn, modelname = modelname)
        assert_in(modelname.encode(), resp, f"Model name '{modelname}' not returned / correctly parsed")
        assert_in(solidname.encode(), resp, f"Solid name '{modelname}' not returned / correctly parsed")
        assert_in(contents, resp, f"STL File contents not returned / correctly parsed")
        self.closeconn(conn)

    def openconn(self):
        self.debug("Connecting to service")
        conn = self.connect()
        conn.read_until("$") # ignore welcome
        return conn

    def closeconn(self, conn: SimpleSocket):
        self.debug("Sending exit command")
        conn.write("exit\n")
        conn.close()

    def putflag(self):  # type: () -> None
        if self.variant_id == 0:
            conn = self.openconn()
            modelname = self.fake_filename()
            stlfile, fileid = self.putfile(conn, solidname = self.flag, modelname = modelname, filetype = "ascii")
            self.closeconn(conn)
            self.postdb({ "fileid": fileid, "modelname": modelname })
        else:
            raise EnoException("Invalid variant_id provided")

    def getflag(self):  # type: () -> None
        if self.variant_id == 0:
            fileid, modelname = self.querydb("fileid", "modelname")
            conn = self.openconn()
            resp = self.getfile(conn, modelname)
            assert_in(self.flag.encode(), resp, "Resulting flag was found to be incorrect")
            self.closeconn(conn)
        else:
            raise EnoException("Invalid variant_id provided")


    def putnoise(self):  # type: () -> None
        if self.variant_id == 0:
            conn = self.openconn()
            modelname = self.fake_filename()
            solidname = self.fake_filename()
            contents, fileid = self.putfile(conn, modelname = modelname, solidname = solidname, filetype = "bin")
            self.closeconn(conn)
            self.postdb({ "fileid": fileid, "modelname": modelname, "solidname": solidname, "contents": contents })
        else:
            raise EnoException("Invalid variant_id provided")

    def getnoise(self):  # type: () -> None
        if self.variant_id == 0:
            fileid, modelname, solidname, contents = self.querydb("fileid", "modelname", "solidname", "contents")
            conn = self.openconn()
            resp = self.getfile(conn, modelname)
            # assert_in(contents, resp, "File content returned by service found to be incorrect")
            assert_in(solidname.encode(), resp, "Solid name returned by service found to be incorrect")
            assert_in(modelname.encode(), resp, "Model name returned by service found to be incorrect")
            self.closeconn(conn)
        else:
            raise EnoException("Invalid variant_id provided")

    def havoc(self):  # type: () -> None
        if self.variant_id == 0:
            self.havoc_upload(filetype = 'ascii', registered = False)
        elif self.variant_id == 1:
            self.havoc_upload(filetype = 'bin', registered = False)
        elif self.variant_id == 2:
            self.havoc_upload(filetype = 'ascii', registered = True)
        elif self.variant_id == 3:
            self.havoc_upload(filetype = 'bin', registered = True)
        else:
            raise EnoException("Invalid variant_id provided");

        # TODO!
        conn = self.openconn()
        self.closeconn(conn)

    def exploit(self):
        """
        This method was added for CI purposes for exploits to be tested.
        Will (hopefully) not be called during actual CTF.
        :raises EnoException on Error
        :return This function can return a result if it wants
                If nothing is returned, the service status is considered okay.
                The preferred way to report Errors in the service is by raising an appropriate EnoException
        """
        # TODO: We still haven't decided if we want to use this function or not. TBA
        pass


app = STLDoctorChecker.service  # This can be used for uswgi.
if __name__ == "__main__":
    run(STLDoctorChecker)