Respond to challenge, if server asks for itHEADmaster

1 files changed, 17 insertions, 3 deletions

diff --git a/QueryServer.py b/QueryServer.py
index 230bbee..7f87178 100755
--- a/QueryServer.py
+++ b/QueryServer.py
@@ -84,12 +84,24 @@ class QueryServer:
 
         message = None
         try:
+            # \x54 A2S_INFO
+            # \x55 A2S_PLAYER
+            # \x56 A2S_RULES
+            # \x56 A2S_RULES
             conn.connect((ip, port))
+            # Fetch A2S_INFO
             conn.sendall(b"\xFF\xFF\xFF\xFF\x54" + b"Source Engine Query" + b"\x00")
             message = conn.recv(4096)
+            if message[4] == 0x41:
+                # If server responds with challenge, return it back
+                conn.sendall(b"\xFF\xFF\xFF\xFF\x54" + b"Source Engine Query" + b"\x00" + message[5:])
+                message = conn.recv(4096)
+
+            # Request Player Info
             conn.sendall(b"\xFF\xFF\xFF\xFF\x55\xFF\xFF\xFF\xFF")
-            challenge = conn.recv(1024)[5:]
-            conn.sendall(b"\xFF\xFF\xFF\xFF\x55" + challenge)
+            challenge = conn.recv(1024)
+            # Answer challenge / retrieve player info
+            conn.sendall(b"\xFF\xFF\xFF\xFF\x55" + challenge[5:])
             playerdata = conn.recv(8192)
         except Exception as e:
             print("Exception in SourceQuery connection:", e)
@@ -103,7 +115,9 @@ class QueryServer:
         latency = time.time() - starttime
         data["latency"] = int(latency * 1000)
 
-        # [0]-[3] -> 4*\xFF
+        # [0]-[3] -> 4*\xFF == Package not split
+        # Otherwise we need to stitch them together
+        assert message[0:4] == b'\xFF\xFF\xFF\xFF'
         data["header"] = chr(message[4])
         data["protocol"] = chr(message[5])
         message = message[6:]