refactor script into remote challenge client

programming/tcp-back-to-school/main.py

@@ -1,28 +1,37 @@
+import math
+import re
 import socket
-from threading import Thread
 
 
-serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-host = "challenge01.root-me.org"
-port = 52002
-serversocket.bind((host, port))
+HOST = "challenge01.root-me.org"
+PORT = 52002
 
 
-class client(Thread):
-    def __init__(self, socket, address):
-        Thread.__init__(self)
-        self.sock = socket
-        self.addr = address
-        self.start()
+def solve_prompt(prompt: str) -> str:
+    match = re.search(
+        r"square root of\s*(\d+)\s*and multiply by\s*(\d+)",
+        prompt,
+        re.IGNORECASE,
+    )
+    if not match:
+        raise ValueError("Unsupported challenge prompt")
 
-    def run(self):
-        while 1:
-            print("Client sent:", self.sock.recv(1024).decode())
-            self.sock.send(b"Oi you sent something to me")
+    left = int(match.group(1))
+    right = int(match.group(2))
+    return str(math.sqrt(left) * right)
 
 
-serversocket.listen(5)
-print("server started and listening")
-while 1:
-    clientsocket, address = serversocket.accept()
-    client(clientsocket, address)
+def main() -> None:
+    with socket.create_connection((HOST, PORT), timeout=2) as conn:
+        prompt = conn.recv(4096).decode(errors="replace")
+        print(prompt, end="")
+
+        answer = solve_prompt(prompt)
+        conn.sendall((answer + "\n").encode())
+
+        result = conn.recv(4096).decode(errors="replace")
+        print(result, end="")
+
+
+if __name__ == "__main__":
+    main()