refactored checker a little

1 files changed, 10 insertions, 12 deletions

diff --git a/checker/src/checker.py b/checker/src/checker.py
index 6b1787e..23f1bfd 100644
--- a/checker/src/checker.py
+++ b/checker/src/checker.py
@@ -45,11 +45,11 @@ class STLDoctorChecker(BaseChecker):
         idlen = random.randint(10, 60)
         return "".join([chr(random.randint(32, 127)) for i in range(idlen)])
 
-    def do_auth(self, conn: SimpleSocket, authstr: str):
+    def do_auth(self, conn, authstr):
         self.write(conn, f"auth {authstr}\n")
         resp = conn.recvuntil("$")
-        print(resp)
-        assert_in(b"Success!", resp, f"Login with pass '{authstr}' failed!");
+        authstr = ensure_bytes(authstr)
+        assert_in(b"Success!", resp, f"Login with pass {authstr} failed!");
 
     def check_listed(self, conn, modelid):
         modelid = ensure_bytes(modelid)
@@ -57,7 +57,7 @@ class STLDoctorChecker(BaseChecker):
         resp = conn.recvuntil("$")
         assert_in(modelid, resp, f"Uploaded model is missing from list command")
 
-    def generate_ascii_file(self, solidname):
+    def genfile_ascii(self, solidname):
         solidname = ensure_bytes(solidname)
 
         if len(solidname) != 0:
@@ -83,7 +83,7 @@ class STLDoctorChecker(BaseChecker):
 
         return content
 
-    def generate_bin_file(self, solidname):
+    def genfile_bin(self, solidname):
         solidname = ensure_bytes(solidname)
 
         if len(solidname) > 78:
@@ -105,11 +105,11 @@ class STLDoctorChecker(BaseChecker):
             content += b"\x00\x00"
         return content
 
-    def generate_file(self, filetype, solidname):
+    def genfile(self, filetype, solidname):
         if filetype == "ascii":
-            return self.generate_ascii_file(solidname = solidname)
+            return self.genfile_ascii(solidname)
         elif filetype == "bin":
-            return self.generate_bin_file(solidname = solidname)
+            return self.genfile_bin(solidname)
         else:
             raise EnoException("Invalid file type supplied");
 
@@ -118,7 +118,7 @@ class STLDoctorChecker(BaseChecker):
         modelname = ensure_bytes(modelname)
 
         # Generate file contents
-        stlfile = self.generate_file(filetype = filetype, solidname = solidname)
+        stlfile = self.genfile(filetype = filetype, solidname = solidname)
 
         # Upload file
         self.debug("Sending command to submit file")
@@ -310,7 +310,7 @@ class STLDoctorChecker(BaseChecker):
         else:
             raise EnoException("Invalid variant_id provided");
 
-    def exploit(self):
+    def exploit(self):  # type: () -> None
         if self.variant_id == 0:
             pass
         elif self.variant_id == 1:
@@ -318,8 +318,6 @@ class STLDoctorChecker(BaseChecker):
         else:
             raise EnoException("Invalid variant_id provided")
 
-        pass
-
 
 app = STLDoctorChecker.service  # This can be used for uswgi.
 if __name__ == "__main__":