141 lines
4.3 KiB
Python
141 lines
4.3 KiB
Python
"""
|
|
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)
|