feat: add timer

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

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

@@ -1,6 +1,7 @@
 import math
 import re
 import socket
+import time
 
 
 HOST = "challenge01.root-me.org"
@@ -22,7 +23,9 @@ def solve_prompt(prompt: str) -> str:
 
 
 def main() -> None:
-    with socket.create_connection((HOST, PORT), timeout=2) as conn:
+    conn = socket.create_connection((HOST, PORT), timeout=2)
+    start_time = time.monotonic()
+    try:
         prompt = conn.recv(4096).decode(errors="replace")
         print(prompt, end="")
 
@@ -31,6 +34,10 @@ def main() -> None:
 
         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__":

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

programming/tcp-back-to-school/solve.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python3 ./main.py | grep "Here is your flag" | cut -d":" -f2 | xargs

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

@@ -0,0 +1,51 @@
+import io
+import unittest
+from contextlib import redirect_stdout
+from unittest.mock import patch
+
+import main
+
+
+class FakeConn:
+    def __init__(self) -> None:
+        self.closed = False
+        self.sent = []
+        self._responses = [
+            b"Calculate the square root of 4 and multiply by 5 = ",
+            b"[+] Good job ! Here is your flag: RM{test}",
+        ]
+
+    def recv(self, _size: int) -> bytes:
+        return self._responses.pop(0)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, _exc_type, _exc, _tb) -> None:
+        self.close()
+
+    def sendall(self, payload: bytes) -> None:
+        self.sent.append(payload)
+
+    def close(self) -> None:
+        self.closed = True
+
+
+class MainTests(unittest.TestCase):
+    def test_main_prints_connection_timer_after_closing(self) -> None:
+        conn = FakeConn()
+        out = io.StringIO()
+
+        with (
+            patch("main.socket.create_connection", return_value=conn),
+            redirect_stdout(out),
+        ):
+            main.main()
+
+        output = out.getvalue()
+        self.assertIn("Connection lifetime:", output)
+        self.assertTrue(conn.closed)
+
+
+if __name__ == "__main__":
+    unittest.main()