diff options
Diffstat (limited to 'checker/src')
| -rw-r--r-- | checker/src/checker.py | 45 |
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: |