def stage_a(self, port: int = START_PORT) -> Packet:
packet = Packet(payload=b'hello world\0',
p_secret=0,
step=self.step,
student_id=self.student_id)
logger.info(
f"[Stage A] Sending packet {packet} to {self.ip_addr}:{port}")
self.udp_socket.sendto(packet.bytes, (self.ip_addr, port))
packet = Packet.from_raw(self.udp_socket.recv(1024))
secret = struct.unpack("!I", packet.payload[-4:])[0]
self.secrets['a'] = secret
logger.info("[Stage A] Finished.")
return packet
def handle_stage_c(self, handler: HookedHandler, packet: Packet):
assert self.secrets[packet.p_secret]["prev_stage"] == "b"
sock = handler.request
num2 = random.randint(1, 10)
len2 = random.randint(1, 10) * 4
secret_c = self.generate_secret()
char = secrets.token_bytes(1)
self.secrets[secret_c] = {
"prev_stage": "c",
"num2": num2,
"len2": len2,
"char": char
}
payload = struct.pack("!3I4s", num2, len2, secret_c, char * 4)
response = Packet(payload=payload,
p_secret=packet.p_secret,
step=2,
student_id=packet.student_id)
logger.info(
f"[Stage C] Sending packet {response} "
f"to {handler.client_address[0]}:{handler.client_address[1]}")
sock.sendall(response.bytes)
self.handle_stage_d(handler)
def stage_d(self, response: Packet):
num2, len2, secret_c, char = struct.unpack("!3I4s", response.payload)
packet = Packet(payload=bytes([char[0]]) * len2,
p_secret=secret_c,
step=self.step,
student_id=self.student_id)
logger.info(
f"[Stage D] Sending {num2} packets with data {packet.bytes} "
f"to {self.ip_addr}:{self.tcp_port}")
self.tcp_socket.sendall(packet.bytes * num2)
packet = Packet.from_raw(self.tcp_socket.recv(1024))
secret = struct.unpack("!I", packet.payload[-4:])[0]
self.secrets['d'] = secret
logger.info("[Stage D] Finished.")
def stage_b(self, response: Packet) -> Packet:
num, length, udp_port, secret_a = struct.unpack(
"!4I", response.payload)
# For stage b, unacknowledged packets should be re-sent
# after 0.5 seconds.
self.udp_socket.settimeout(0.5)
for packet_id in range(num):
while True:
payload = struct.pack(f"!I{length}s", packet_id,
b'\0' * length)
packet = Packet(payload=payload,
p_secret=secret_a,
step=self.step,
student_id=self.student_id)
logger.info(
f"[Stage B] Sending packet {packet} to {self.ip_addr}:{udp_port}"
)
self.udp_socket.sendto(packet.bytes, (self.ip_addr, udp_port))
try:
response_packet = Packet.from_raw(
self.udp_socket.recv(1024))
ack = struct.unpack("!I", response_packet.payload)[0]
if ack == packet_id:
logger.info(
f"[Stage B] Packet acknowledged (id: {packet_id})")
break
else:
logger.info(
f"[Stage B] Packet (id: {packet_id}) not acknowledged: "
f"expected payload {packet_id}, got {ack}. Retrying."
)
except socket.timeout:
logger.info(
f"[Stage B] Packet dropped (id: {packet_id}). Retrying."
)
self.udp_socket.settimeout(None)
packet = Packet.from_raw(self.udp_socket.recv(1024))
secret = struct.unpack("!I", packet.payload[-4:])[0]
self.secrets['b'] = secret
logger.info("[Stage B] Finished.")
return packet
def handle_stage_d(self, handler: HookedHandler):
sock = handler.request
data = sock.recv(2048)
logger.info(f"[Stage D] Received data {data} from "
f"{handler.client_address[0]}:{handler.client_address[1]}")
payload_len, secret_c, step, student_id = struct.unpack(
"!IIHH", data[:12])
try:
prev_stage = self.secrets[secret_c]['prev_stage']
num2 = self.secrets[secret_c]["num2"]
len2 = self.secrets[secret_c]["len2"]
char = self.secrets[secret_c]["char"]
except KeyError:
logger.error(f"Unrecognized secret: {secret_c}")
return
assert prev_stage == "c"
for i in range(num2):
start, end = i * (len2 + 12), (i + 1) * (len2 + 12)
try:
packet = Packet.from_raw(data[start:end])
except ValueError as e:
# Malformed packet
logger.info(e)
return
if packet.payload != char * len2:
logger.error(
f"[Stage D] Packet does not conform to protocol. "
f"Packet info: {packet!r}\n"
f"packet.payload: {packet.payload}, expected: {char * len2}"
)
return
payload = struct.pack("!I", self.generate_secret())
response = Packet(payload=payload,
p_secret=secret_c,
step=2,
student_id=student_id)
logger.info(
f"[Stage D] Sending packet {response} "
f"to {handler.client_address[0]}:{handler.client_address[1]}")
sock.sendall(response.bytes)
def stage_c(self, response: Packet) -> Packet:
tcp_port, secret_b = struct.unpack("!II", response.payload)
self.tcp_socket.connect((self.ip_addr, tcp_port))
self.tcp_port = tcp_port
logger.info(
f"[Stage C] Connected to TCP socket at {self.ip_addr}:{tcp_port}")
packet = Packet.from_raw(self.tcp_socket.recv(1024))
secret = struct.unpack("!I", packet.payload[-4:])[0]
self.secrets['c'] = secret
logger.info("[Stage C] Finished.")
return packet
def handle_stage_b(self, handler: HookedHandler):
if handler.server.fileno() == -1:
client_ip, client_port = handler.client_address
server_ip, server_port = handler.server.server_address
logger.error(
f"[Stage B] Received data from {client_ip}:{client_port} "
f"but the server at {server_ip}:{server_port} is already closed."
)
return
data, sock = handler.request
logger.info(f"[Stage B] Received packet {data} from "
f"{handler.client_address[0]}:{handler.client_address[1]}")
try:
packet = Packet.from_raw(data)
except ValueError as e:
# Malformed packet
logger.error(e)
return
try:
prev_stage = self.secrets[packet.p_secret]['prev_stage']
num_packets = self.secrets[packet.p_secret]["num_packets"]
remaining_packets = self.secrets[
packet.p_secret]["remaining_packets"]
ack_fails = self.secrets[packet.p_secret]["ack_fails"]
packet_len = self.secrets[packet.p_secret]['packet_len']
packet_id = struct.unpack("!I", packet.payload[:4])[0]
except KeyError:
logger.error(f"Unrecognized secret: {packet.p_secret}")
return
except struct.error:
logger.error("Packet payload must be at least 4 bytes long")
return
assert prev_stage == "a"
# Ensure that this packet is a valid packet for step b1
if (packet.step != 1 or packet_len + 4 != len(packet.payload)):
logger.error(
f"[Stage B] Packet does not conform to protocol. "
f"Packet info: {packet!r}\n"
f"packet.step: {packet.step}, expected: 1\n"
f"len(packet.payload): {len(packet.payload)}, "
f"expected: {packet_len + 4}\n"
f"packet_id: {packet_id}, expected: {num_packets - remaining_packets}, "
f"num_packets: {num_packets}, remaining_packets: {remaining_packets}"
)
return
if packet_id in ack_fails:
logger.info(f"[Stage B] Dropping packet with id {packet_id}")
ack_fails.remove(packet_id)
return
else:
logger.info(f"[Stage B] Acknowledging packet with id {packet_id}")
if num_packets - remaining_packets == packet_id:
# Only decrement remaining_packets if this is not a resent packet
self.secrets[packet.p_secret]["remaining_packets"] -= 1
ack = Packet(payload=packet.payload[:4],
p_secret=packet.p_secret,
step=2,
student_id=packet.student_id)
sock.sendto(ack.bytes, handler.client_address)
if self.secrets[packet.p_secret]["remaining_packets"] == 0:
secret_b = self.generate_secret()
self.secrets[secret_b] = {'prev_stage': "b"}
tcp_port = self.random_port()
payload = struct.pack("!II", tcp_port, secret_b)
response = Packet(payload=payload,
p_secret=secret_b,
step=2,
student_id=packet.student_id)
server = TimeoutThreadingTCPServer(
(self.ip_addr, tcp_port),
self.handler_factory(callback=self.handle_stage_c,
callback_args=(response, )),
timeout=TIMEOUT,
after_close=synchronized(
lambda: self.tcp_servers.pop(tcp_port), lock=self.rlock))
assert tcp_port not in self.tcp_servers
self.tcp_servers[tcp_port] = server
self.start_server(server)
logger.info(
f"[Stage B] Started new TCP server at {self.ip_addr}:{tcp_port}."
)
logger.info(
f"[Stage B] Sending packet {response} "
f"to {handler.client_address[0]}:{handler.client_address[1]}")
sock.sendto(response.bytes, handler.client_address)
def handle_stage_a(self, handler: HookedHandler):
data, sock = handler.request
logger.info(f"[Stage A] Received packet {data} from "
f"{handler.client_address[0]}:{handler.client_address[1]}")
try:
packet = Packet.from_raw(data)
except ValueError as e:
# Packet is malformed
logger.error(e)
return
if (packet.payload.lower() != b'hello world\0' or packet.p_secret != 0
or packet.step != 1 or packet.payload_len != 12):
logger.error(f"[Stage A] Packet does not conform to protocol. "
f"Packet info: {packet!r}")
return
num_packets = random.randint(5, 10)
# Guarantee that packet_len is a multiple of 4.
packet_len = random.randint(1, 10) * 4
secret_a = self.generate_secret()
udp_port = self.random_port()
server = TimeoutThreadingUDPServer(
(self.ip_addr, udp_port),
self.handler_factory(callback=self.handle_stage_b),
timeout=TIMEOUT,
after_close=synchronized(lambda: self.udp_servers.pop(udp_port),
lock=self.rlock))
assert udp_port not in self.udp_servers
self.udp_servers[udp_port] = server
self.start_server(server)
logger.info(
f"[Stage A] Started new UDP server at {self.ip_addr}:{udp_port}.")
# For stage B, don't acknowledge one packet
# ack_fails = set()
ack_fails = {random.choice(range(num_packets))}
# ack_fails = set(random.sample(
# range(num_packets),
# k=random.randint(1, num_packets))
# )
# Store relevant data for stage B
self.secrets[secret_a] = {
'prev_stage': "a",
"num_packets": num_packets,
# Number of packets that the server is still expecting to receive
"remaining_packets": num_packets,
# Packets that we should drop once
"ack_fails": ack_fails,
"packet_len": packet_len
}
payload = struct.pack("!4I", num_packets, packet_len, udp_port,
secret_a)
response = Packet(payload=payload,
p_secret=0,
step=2,
student_id=packet.student_id)
logger.info(
f"[Stage A] Sending packet {response} "
f"to {handler.client_address[0]}:{handler.client_address[1]}")
sock.sendto(response.bytes, handler.client_address)