#!/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, selectors, time, socket
import numpy as np

logging.getLogger("faker").setLevel(logging.WARNING)
logging.getLogger("pwnlib").setLevel(logging.WARNING)

from faker import Faker

def ensure_bytes(v):
    if type(v) == bytes:
        return v
    elif type(v) == str:
        return v.encode("latin-1")
    else:
        raise BrokenServiceException("Tried to pass non str/bytes to bytes arg")

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

    debuglog = True

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

    def write(self, conn, buf):
        if self.debuglog:
            print("SEND: " + str(ensure_bytes(buf)))
        conn.send(ensure_bytes(buf))

    def fakeid(self):
        fake = Faker(["en_US"])
        allowed = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmopqrstuvwxyz0123456789-+.!"
        return "".join([c for c in fake.name().replace(' ','') if c in allowed][:60]).ljust(10, '.')

    def havocid(self):
        idlen = random.randint(10, 60)
        return "".join([chr(random.randint(32, 127)) for i in range(idlen)])

    def do_auth(self, conn: SimpleSocket, authstr: str):
        self.write(conn, f"auth {authstr}\n")
        resp = conn.recvuntil("$")
        print(resp)
        assert_in(b"Success!", resp, f"Login with pass '{authstr}' failed!");

    def check_listed(self, conn, modelid):
        modelid = ensure_bytes(modelid)
        self.write(conn, "list\n")
        resp = conn.recvuntil("$")
        assert_in(modelid, resp, f"Uploaded model is missing from list command")

    def generate_ascii_file(self, solidname):
        solidname = ensure_bytes(solidname)

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

        return content

    def generate_bin_file(self, solidname):
        solidname = ensure_bytes(solidname)

        if len(solidname) > 78:
            raise EnoException("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
        facet_count = random.randint(4, 30)
        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, solidname):
        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, modelname, solidname, filetype):
        solidname = ensure_bytes(solidname)
        modelname = ensure_bytes(modelname)

        # Generate file contents
        stlfile = self.generate_file(filetype = filetype, solidname = solidname)

        # Upload file
        self.debug("Sending command to submit file")
        self.write(conn, "upload\n")
        self.write(conn, f"{len(stlfile)}\n")
        self.write(conn, stlfile)
        self.write(conn, modelname + b"\n")

        # Parse ID
        print(conn.recvuntil("with ID "))
        modelid = conn.recvuntil(b"!")[:-1]
        if modelid == b"":
            raise BrokenServiceException("Unable to upload file!")
        self.debug(f"Got ID {modelid}")

        conn.recvuntil(b"$")

        if self.debuglog:
            print(f"PUT FILE: {modelid}")

        return stlfile, modelid

    def getfile(self, conn, modelname):
        modelname = ensure_bytes(modelname)

        if modelname != "":
            self.debug(f"Sending command to retrieve file with name '{modelname}'")
            self.write(conn, "search\n")
            self.write(conn, modelname + b"\n")
            self.write(conn, "0\n") # first result
            self.write(conn, "y\n") # yes download
        else:
            self.debug(f"Sending command to retrieve file")
            self.write(conn, f"search\n0\ny\n")

        resp = conn.recvuntil(b"Here you go.. (")
        try:
            size = int(conn.recvuntil(b"B)\n")[:-3])
        except:
            raise BrokenServiceException("Returned file content size for download is not a valid integer")

        print("FILE SIZE:", size)
        contents = conn.recvn(size)
        if self.debuglog:
            print("GOT FILE:\n" + str(contents))

        conn.recvuntil("$") # clean up rest
        return resp + contents

    def check_getfile(self, conn, modelname, solidname, contents, modelid = None):
        resp = self.getfile(conn, modelname = modelname)
        if modelid:
            assert_in(ensure_bytes(modelid), resp, f"Model id {modelid} not returned / correctly parsed")
        assert_in(ensure_bytes(modelname), resp, f"Model name {modelname} not returned / correctly parsed")
        assert_in(ensure_bytes(solidname), resp, f"Solid name {solidname} not returned / correctly parsed")
        assert_in(ensure_bytes(contents), resp, f"STL File contents not returned / correctly parsed")

    def querydb(self, *args):
        self.debug("Querying db contents");
        vals = []
        for arg in args:
            try:
                val: str = self.chain_db[arg]
            except KeyError 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, register):
        # cant be havocid with ascii since might mess with stl parsing
        solidname = self.fakeid() if filetype == 'ascii' else self.havocid()
        modelname = self.havocid()
        authstr = self.havocid()

        # create new session and user and upload file
        conn = self.openconn()
        if register:
            self.do_auth(conn, authstr)
        contents, modelid = self.putfile(conn, modelname, solidname, filetype)
        self.check_getfile(conn, modelname, solidname, contents)
        if register:
            self.check_listed(conn, modelid)
        self.closeconn(conn)

        # try getting file from a new session
        conn = self.openconn()
        if register:
            self.do_auth(conn, authstr)
        self.check_getfile(conn, modelname, solidname, contents)
        if register:
            self.check_listed(conn, modelid)
        self.closeconn(conn)

    def openconn(self):
        import pwnlib
        self.debug("Connecting to service")
        conn = pwnlib.tubes.remote.remote(self.address, self.port, timeout = self.timeout)
        conn.recvuntil("$") # ignore welcome
        return conn

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

    def putflag(self):  # type: () -> None
        if self.variant_id == 0:
            conn = self.openconn()
            modelname = self.fakeid()
            stlfile, modelid = self.putfile(conn, modelname, self.flag, filetype = "ascii")
            self.closeconn(conn)
            self.postdb({ "modelid": modelid, "modelname": modelname })
        elif self.variant_id == 1:
            conn = self.openconn()
            modelname = self.fakeid()
            authstr = self.fakeid()
            self.do_auth(conn, authstr)
            stlfile, modelid = self.putfile(conn, modelname, self.flag, filetype = "bin")
            self.closeconn(conn)
            self.postdb({ "modelid": modelid, "modelname": modelname, "auth": authstr })
        else:
            raise EnoException("Invalid variant_id provided")

    def getflag(self):  # type: () -> None
        if self.variant_id == 0:
            modelid, modelname = self.querydb("modelid", "modelname")
            conn = self.openconn()
            resp = self.getfile(conn, modelname.encode())
            self.debug(resp)
            assert_in(self.flag.encode(), resp, "Resulting flag was found to be incorrect")
            self.closeconn(conn)
        elif self.variant_id == 1:
            modelid, modelname, authstr = self.querydb("modelid", "modelname", "auth")
            conn = self.openconn()
            self.do_auth(conn, authstr)
            resp = self.getfile(conn, modelname.encode())
            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.fakeid()
            solidname = self.fakeid()
            contents, modelid = self.putfile(conn, modelname, solidname, "bin")
            self.closeconn(conn)
            self.postdb({ "modelid": modelid, "modelname": modelname, "solidname": solidname, "contents": contents })
        elif self.variant_id == 1:
            conn = self.openconn()
            authstr = self.fakeid()
            modelname = self.fakeid()
            solidname = self.fakeid()
            self.do_auth(conn, authstr)
            contents, modelid = self.putfile(conn, modelname, solidname, "ascii")
            self.closeconn(conn)
            self.postdb({ "modelid": modelid, "modelname": modelname, "solidname": solidname, "contents": contents, "auth": authstr })
        else:
            raise EnoException("Invalid variant_id provided")

    def getnoise(self):  # type: () -> None
        if self.variant_id == 0:
            modelid, modelname, solidname, contents = self.querydb("modelid", "modelname", "solidname", "contents")
            conn = self.openconn()
            self.check_getfile(conn, modelname, solidname, contents, modelid)
            self.closeconn(conn)
        elif self.variant_id == 1:
            modelid, modelname, solidname, contents, authstr = self.querydb("modelid", "modelname", "solidname", "contents", "auth")
            conn = self.openconn()
            self.do_auth(conn, authstr)
            self.check_getfile(conn, modelname, solidname, contents, modelid)
            self.closeconn(conn)
        else:
            raise EnoException("Invalid variant_id provided")

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

    def exploit(self):
        if self.variant_id == 0:
            pass
        elif self.variant_id == 1:
            pass
        else:
            raise EnoException("Invalid variant_id provided")

        pass


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