47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
"""
|
|
RFC 1071 one's complement checksum + TCP/UDP pseudo-header helpers.
|
|
|
|
These are the only "tricky" bits in the stack — we implement them
|
|
deliberately and have unit tests against known IPv4 / TCP examples
|
|
from the RFCs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import struct
|
|
|
|
|
|
def ones_complement_sum(data: bytes) -> int:
|
|
"""RFC 1071 16-bit one's complement sum, returned BEFORE inversion."""
|
|
if len(data) & 1:
|
|
data = data + b'\x00'
|
|
s = 0
|
|
for hi, lo in zip(data[0::2], data[1::2]):
|
|
s += (hi << 8) | lo
|
|
while s >> 16:
|
|
s = (s & 0xffff) + (s >> 16)
|
|
return s
|
|
|
|
|
|
def internet_checksum(data: bytes) -> int:
|
|
"""Final checksum: invert the one's-complement sum."""
|
|
return (~ones_complement_sum(data)) & 0xffff
|
|
|
|
|
|
def tcp_udp_pseudo_header(src_ip: bytes, dst_ip: bytes, protocol: int, length: int) -> bytes:
|
|
"""
|
|
The pseudo-header that goes into TCP/UDP checksum but isn't
|
|
transmitted. Per RFC 793 §3.1:
|
|
|
|
src_ip (4) dst_ip (4) zero (1) protocol (1) length (2)
|
|
"""
|
|
return src_ip + dst_ip + bytes([0, protocol]) + struct.pack('!H', length)
|
|
|
|
|
|
def tcp_udp_checksum(
|
|
src_ip: bytes, dst_ip: bytes, protocol: int, segment: bytes
|
|
) -> int:
|
|
"""Compute the checksum field for a TCP or UDP segment."""
|
|
pseudo = tcp_udp_pseudo_header(src_ip, dst_ip, protocol, len(segment))
|
|
return internet_checksum(pseudo + segment)
|