CTF toolkit: lib (net, crypto_utils), templates per category, scaffold, cheatsheet, ps_and_qs example

lib/net.py

@@ -0,0 +1,140 @@
+"""
+lib/net.py — Connection helpers (p4-team style: nc / receive_until_match / send).
+
+Works with pwntools if installed; otherwise falls back to a raw-socket
+implementation so your solvers run even on a bare Python.
+"""
+import os
+import re
+import socket
+import sys
+import time
+
+try:
+    from pwn import remote, context, p64, u64, ELF, ROP  # noqa
+    HAVE_PWN = True
+except Exception:
+    HAVE_PWN = False
+
+
+class Conn:
+    """Thin wrapper around pwntools.remote, or a raw socket if pwntools is missing."""
+
+    def __init__(self, host, port, timeout=10, use_pwntools=True):
+        self.host = host
+        self.port = port
+        self.timeout = timeout
+        self.use_pwntools = use_pwntools and HAVE_PWN
+        if self.use_pwntools:
+            context.log_level = os.environ.get("CTF_LOG", "info")
+            self.s = remote(host, port, timeout=timeout)
+        else:
+            self.s = socket.create_connection((host, port), timeout=timeout)
+            self._buf = b""
+
+    # ---- low level ----
+    def recv_raw(self, n=4096):
+        if self.use_pwntools:
+            return self.s.recv(n)
+        data = b""
+        try:
+            while len(data) < n:
+                chunk = self.s.recv(n - len(data))
+                if not chunk:
+                    break
+                data += chunk
+        except socket.timeout:
+            pass
+        return data
+
+    def recv_until(self, marker, timeout=None):
+        """Receive until `marker` (bytes) appears. Returns everything including marker."""
+        if self.use_pwntools:
+            return self.s.recvuntil(marker)
+        marker = marker.encode() if isinstance(marker, str) else marker
+        end = time.time() + (timeout or self.timeout)
+        buf = self._buf
+        while marker not in buf and time.time() < end:
+            try:
+                self.s.settimeout(max(0.1, end - time.time()))
+                chunk = self.s.recv(4096)
+                if not chunk:
+                    break
+                buf += chunk
+            except socket.timeout:
+                break
+        self._buf = b""
+        return buf
+
+    def recv_until_match(self, pat, timeout=None):
+        """Receive until regex `pat` matches. Returns the matched prefix+match."""
+        if self.use_pwntools:
+            return self.s.recvline_regex(pat) if hasattr(self.s, "recvline_regex") else self.s.recvuntil(pat.encode())
+        rx = re.compile(pat.encode() if isinstance(pat, str) else pat)
+        end = time.time() + (timeout or self.timeout)
+        buf = self._buf
+        while time.time() < end:
+            m = rx.search(buf)
+            if m:
+                self._buf = buf[m.end():]
+                return buf[:m.end()]
+            try:
+                self.s.settimeout(max(0.1, end - time.time()))
+                chunk = self.s.recv(4096)
+                if not chunk:
+                    break
+                buf += chunk
+            except socket.timeout:
+                break
+        return buf
+
+    def send(self, data):
+        data = data.encode() if isinstance(data, str) else data
+        if self.use_pwntools:
+            self.s.sendline(data) if data.endswith(b"\n") else self.s.send(data)
+        else:
+            self.s.sendall(data)
+
+    def sendline(self, data):
+        data = data.encode() if isinstance(data, str) else data
+        if self.use_pwntools:
+            self.s.sendline(data)
+        else:
+            self.s.sendall(data + b"\n")
+
+    def interactive(self):
+        if self.use_pwntools:
+            self.s.interactive()
+        else:
+            print("[!] interactive needs pwntools; dropping to manual mode")
+            import select
+            while True:
+                r, _, _ = select.select([self.s, sys.stdin], [], [])
+                if self.s in r:
+                    print(self.recv_raw(4096).decode(errors="replace"), end="")
+                if sys.stdin in r:
+                    self.s.sendall(sys.stdin.readline().encode())
+
+    def close(self):
+        try:
+            self.s.close()
+        except Exception:
+            pass
+
+
+def nc(host, port, timeout=10):
+    """Drop-in for the p4 pattern `nc(host, port)`."""
+    return Conn(host, port, timeout=timeout)
+
+
+# convenience re-exports for templates
+def receive_until(s, marker, timeout=None):
+    return s.recv_until(marker, timeout)
+
+
+def receive_until_match(s, pat, timeout=None):
+    return s.recv_until_match(pat, timeout)
+
+
+def send(s, data):
+    s.send(data)