feat: add timer

Signed-off-by: Tuan-Dat Tran <tuan-dat.tran@dextradata.com>

programming/tcp-back-to-school/notes.org

@@ -24,32 +24,63 @@ cat ./main.py
 
 #+RESULTS:
 #+begin_example
+import math
+import re
 import socket
-from threading import Thread
+import time
 
 
-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 f"{math.sqrt(left) * right:.2f}"
 
 
-serversocket.listen(5)
-print("server started and listening")
-while 1:
-    clientsocket, address = serversocket.accept()
-    client(clientsocket, address)
+def main() -> None:
+    conn = socket.create_connection((HOST, PORT), timeout=2)
+    start_time = time.monotonic()
+    try:
+        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="")
+    finally:
+        conn.close()
+        elapsed = time.monotonic() - start_time
+        print(f"\nConnection lifetime: {elapsed:.3f}s")
+
+
+if __name__ == "__main__":
+    main()
 #+end_example
+
+#+begin_src sh :results output
+python3 ./main.py
+#+end_src
+
+#+RESULTS:
+: 
+: ====================
+:  GO BACK TO COLLEGE
+: ====================
+: You should tell me the answer of this math operation in less than 2 seconds !
+: 
+: Calculate the square root of 535 and multiply by 2447 = [+] Good job ! Here is your flag: RM{TCP_C0nnecT_4nD_m4Th}
+: 
+: Connection lifetime: 0.408s