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

#### Checker Tenets
# A checker SHOULD not be easily identified by the examination of network traffic
#    => satisfied, because checker uses regular user interface and picks strings from a wordlist
#       to appear more human (TODO)
# A checker SHOULD use unusual, incorrect or pseudomalicious input to detect network filters
#    => satisfied, send various garbage bytes for model name and file contents (TODO)
####

samplestl = """
solid
	facet normal 1.0 0 0
		outer loop
			vertex 0 1 0
			vertex 0 1 1
			vertex 0 0 1
		endloop
	endfacet
	facet normal 0 0 1.0
		outer loop
			vertex 1 0 0
			vertex 1 1 0
			vertex 0 1 0
		endloop
	endfacet
endsolid
"""

class STLDoctorChecker(BaseChecker):
    """
    Change the methods given here, then simply create the class and .run() it.
    Magic.
    A few convenient methods and helpers are provided in the BaseChecker.
    ensure_bytes and ensure_unicode to make sure strings are always equal.
    As well as methods:
    self.connect() connects to the remote server.
    self.get and self.post request from http.
    self.chain_db is a dict that stores its contents to a mongodb or filesystem.
    conn.readline_expect(): fails if it's not read correctly
    To read the whole docu and find more goodies, run python -m pydoc enochecker
    (Or read the source, Luke)
    """

    ##### EDIT YOUR CHECKER PARAMETERS
    flag_variants = 1
    noise_variants = 1
    havoc_variants = 0
    service_name = "stldoctor"
    port = 9000
    ##### END CHECKER PARAMETERS

    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 generate_file(self, filetype: str = "ascii", extra: str = ""):
        if filetype == "ascii":
            # TODO handle extra as solidname and gen randomly
            return samplestl
        elif filetype == "bin":
            # TODO handle extra as header
            return samplestl # TODO: this is not a binary STL!
        else:
            raise EnoException("Invalid file type supplied");

    def putfile(self, conn: SimpleSocket, solidname = "TODO", modelname = "TODO"):
        # Generate file contents
        stlfile = samplestl

        # Upload file
        self.debug("Sending command to submit file")
        conn.write(f"submit\n{len(stlfile)}\n{stlfile}{modelname}\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 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
        """
        This method stores a flag in the service.
        In case multiple flags are provided, self.variant_id gives the appropriate index.
        The flag itself can be retrieved from self.flag.
        On error, raise an Eno Exception.
        :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
        """
        if self.variant_id == 0:
            conn = self.openconn()
            modelname = self.flag
            stlfile, fileid = self.putfile(conn, modelname = modelname)
            self.closeconn(conn)
            self.chain_db = { "fileid": fileid, "modelname": modelname }
        else:
            raise EnoException("Wrong variant_id provided")

    def getflag(self):  # type: () -> None
        """
        This method retrieves a flag from the service.
        Use self.flag to get the flag that needs to be recovered and self.round to get the round the flag was placed in.
        On error, raise an EnoException.
        :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
        """
        if self.variant_id == 0:
            fileid, modelname = self.querydb("fileid", "modelname")
            conn = self.openconn()
            resp = self.getfile(conn, modelname)
            assert_in(modelname.encode(), resp, "Resulting flag was found to be incorrect")
            self.closeconn(conn)
        else:
            raise EnoException("Wrong variant_id provided")


    def putnoise(self):  # type: () -> None
        """
        This method stores noise in the service. The noise should later be recoverable.
        The difference between noise and flag is, that noise does not have to remain secret for other teams.
        This method can be called many times per round. Check how often using self.variant_id.
        On error, raise an EnoException.
        :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
        """
        if self.variant_id == 0:
            conn = self.openconn()
            modelname = "NOISE" # TODO
            contents, fileid = self.putfile(conn, modelname = modelname)
            self.closeconn(conn)
            self.postdb({ "fileid": fileid, "modelname": modelname, "contents": contents })
        else:
            raise EnoException("Wrong variant_id provided")

    def getnoise(self):  # type: () -> None
        """
        This method retrieves noise in the service.
        The noise to be retrieved is inside self.flag
        The difference between noise and flag is, that noise does not have to remain secret for other teams.
        This method can be called many times per round. Check how often using variant_id.
        On error, raise an EnoException.
        :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
        """
        if self.variant_id == 0:
            fileid, modelname, contents = self.querydb("fileid", "modelname", "contents")
            conn = self.openconn()
            resp = self.getfile(conn, modelname)
            assert_in(contents.encode(), resp, "Noise file content was found to be incorrect")
            self.closeconn(conn)
        else:
            raise EnoException("Wrong variant_id provided")

    def havoc(self):  # type: () -> None
        """
        This method unleashes havoc on the app -> Do whatever you must to prove the service still works. Or not.
        On error, raise an EnoException.
        :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
        """
        return

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