added ability to request mulitple files in search without restarting

1 files changed, 19 insertions, 26 deletions

diff --git a/checker/src/checker.py b/checker/src/checker.py
index f65a033..69c9084 100644
--- a/checker/src/checker.py
+++ b/checker/src/checker.py
@@ -204,6 +204,7 @@ class STLDoctorChecker(BaseChecker):
         conn.write(modelname + b"\n")
         conn.write("0\n") # first result
         conn.write("y\n" if download else "\n")
+        conn.write("q\n") # quit
 
         # Wait for end of info box
         resp = conn.recvuntil("================== \n")
@@ -353,39 +354,31 @@ class STLDoctorChecker(BaseChecker):
             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]
+            filelist = [l.strip().split(b" : ")[1] for l in conn.recvuntil("? ").split(b"\n") if b" : " in l]
             if len(filelist) == 0:
                 raise BrokenServiceException("Failed to list files through search")
-            index_dict = {fl[1]: fl[0] for fl in filelist}
-            targets = [fl[1] for fl in filelist]
 
             # Use it to enumerate other files and grab contents
             found = None
-            self.debug("Targets:\n" + "\n".join([' - ' + l.decode('latin1') for l in targets]))
-            for i,fhash in enumerate(targets):
-                if index_dict[fhash] == None:
-                    self.debug(b"Skipping now missing file " + fhash)
-                    continue
-
-                # Retrieve current file
-                self.debug(b"Retrieving file " + fhash + b" at index " + index_dict[fhash])
-                conn.write(index_dict[fhash] + b"\ny\n")
-                fileinfo = conn.recvuntil(self.prompt)
-                #self.debug("File contents:\n" + fileinfo.decode("latin1"))
-                found = self.search_flag_bytes(fileinfo)
-                if found is not None or i == len(targets) - 1:
+            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}\ny\n")
+                resp = conn.recvuntil("==================")
+                resp += conn.recvuntil(b"Here you go.. (")
+                try:
+                    size = int(conn.recvuntil(b"B)\n")[:-3])
+                except:
+                    raise BrokenServiceException("Download size is not a valid integer")
+                resp += conn.recvn(size)
+                resp += conn.recvuntil("? ")
+                self.debug(resp)
+                found = self.search_flag_bytes(resp)
+                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")
-
-                # Update indicies from new search
-                filelist = [l.strip().split(b" : ") for l in conn.recvuntil("?").split(b"\n") if b" : " in l]
-                index_dict = {name : None for name in targets}
-                for fl in filelist:
-                    index_dict[fl[1]] = fl[0]
-
+            conn.write("q\n")
+            conn.recvuntil(self.prompt)
             self.closeconn(conn)
 
             if found is None: