534 lines
17 KiB
Python
534 lines
17 KiB
Python
"""
|
|
Layer-2/3/4 protocol parsers and encoders.
|
|
|
|
Pure-data dataclasses with `from_bytes` / `to_bytes`. No I/O, no async,
|
|
no network state — just bytes ↔ structs. This file is the reference
|
|
that every other module in this package depends on.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import struct
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
from .checksums import internet_checksum, tcp_udp_checksum
|
|
from .consts import (
|
|
BROADCAST_MAC,
|
|
ETHERTYPE_ARP,
|
|
ETHERTYPE_IPV4,
|
|
IPPROTO_ICMP,
|
|
IPPROTO_TCP,
|
|
IPPROTO_UDP,
|
|
)
|
|
|
|
|
|
# ─── Ethernet ────────────────────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class Ethernet:
|
|
dst: bytes
|
|
src: bytes
|
|
ethertype: int
|
|
payload: bytes
|
|
|
|
@classmethod
|
|
def parse(cls, frame: bytes) -> 'Ethernet':
|
|
if len(frame) < 14:
|
|
raise ValueError(f'ethernet frame too short: {len(frame)}')
|
|
return cls(
|
|
dst=frame[0:6],
|
|
src=frame[6:12],
|
|
ethertype=(frame[12] << 8) | frame[13],
|
|
payload=frame[14:],
|
|
)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
return self.dst + self.src + struct.pack('!H', self.ethertype) + self.payload
|
|
|
|
|
|
# ─── ARP ─────────────────────────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class Arp:
|
|
htype: int = 1 # Ethernet
|
|
ptype: int = ETHERTYPE_IPV4
|
|
hlen: int = 6
|
|
plen: int = 4
|
|
opcode: int = 1 # 1=request, 2=reply
|
|
sha: bytes = b'' # sender hw addr
|
|
spa: bytes = b'' # sender protocol (IP) addr
|
|
tha: bytes = b'' # target hw addr
|
|
tpa: bytes = b'' # target protocol addr
|
|
|
|
@classmethod
|
|
def parse(cls, payload: bytes) -> 'Arp':
|
|
if len(payload) < 28:
|
|
raise ValueError(f'ARP payload too short: {len(payload)}')
|
|
htype, ptype, hlen, plen, opcode = struct.unpack('!HHBBH', payload[:8])
|
|
sha = payload[8:14]
|
|
spa = payload[14:18]
|
|
tha = payload[18:24]
|
|
tpa = payload[24:28]
|
|
return cls(htype, ptype, hlen, plen, opcode, sha, spa, tha, tpa)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
return (
|
|
struct.pack('!HHBBH', self.htype, self.ptype, self.hlen, self.plen, self.opcode)
|
|
+ self.sha + self.spa + self.tha + self.tpa
|
|
)
|
|
|
|
|
|
# ─── IPv4 ────────────────────────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class IPv4:
|
|
version: int = 4
|
|
ihl: int = 5 # 5 → 20-byte header (no options)
|
|
dscp: int = 0
|
|
ecn: int = 0
|
|
total_length: int = 0 # populated in to_bytes
|
|
ident: int = 0
|
|
flags: int = 2 # Don't Fragment
|
|
frag_offset: int = 0
|
|
ttl: int = 64
|
|
protocol: int = 0
|
|
checksum: int = 0 # populated in to_bytes
|
|
src: bytes = b'\x00\x00\x00\x00'
|
|
dst: bytes = b'\x00\x00\x00\x00'
|
|
payload: bytes = b''
|
|
|
|
@classmethod
|
|
def parse(cls, data: bytes) -> 'IPv4':
|
|
if len(data) < 20:
|
|
raise ValueError(f'IPv4 too short: {len(data)}')
|
|
b0 = data[0]
|
|
version = b0 >> 4
|
|
ihl = b0 & 0x0f
|
|
if ihl < 5:
|
|
raise ValueError(f'bad IHL {ihl}')
|
|
header_len = ihl * 4
|
|
dscp_ecn = data[1]
|
|
total_length = struct.unpack('!H', data[2:4])[0]
|
|
ident = struct.unpack('!H', data[4:6])[0]
|
|
flags_frag = struct.unpack('!H', data[6:8])[0]
|
|
flags = flags_frag >> 13
|
|
frag_offset = flags_frag & 0x1fff
|
|
ttl = data[8]
|
|
protocol = data[9]
|
|
checksum = struct.unpack('!H', data[10:12])[0]
|
|
src = data[12:16]
|
|
dst = data[16:20]
|
|
payload = data[header_len:total_length]
|
|
return cls(
|
|
version=version,
|
|
ihl=ihl,
|
|
dscp=dscp_ecn >> 2,
|
|
ecn=dscp_ecn & 0x3,
|
|
total_length=total_length,
|
|
ident=ident,
|
|
flags=flags,
|
|
frag_offset=frag_offset,
|
|
ttl=ttl,
|
|
protocol=protocol,
|
|
checksum=checksum,
|
|
src=src,
|
|
dst=dst,
|
|
payload=payload,
|
|
)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
ihl = 5 # we never emit options
|
|
total_length = ihl * 4 + len(self.payload)
|
|
flags_frag = (self.flags << 13) | self.frag_offset
|
|
dscp_ecn = (self.dscp << 2) | self.ecn
|
|
|
|
# Header without checksum first
|
|
header_no_cksum = struct.pack(
|
|
'!BBHHHBBH4s4s',
|
|
(4 << 4) | ihl,
|
|
dscp_ecn,
|
|
total_length,
|
|
self.ident,
|
|
flags_frag,
|
|
self.ttl,
|
|
self.protocol,
|
|
0,
|
|
self.src,
|
|
self.dst,
|
|
)
|
|
cksum = internet_checksum(header_no_cksum)
|
|
header = (
|
|
header_no_cksum[:10]
|
|
+ struct.pack('!H', cksum)
|
|
+ header_no_cksum[12:]
|
|
)
|
|
return header + self.payload
|
|
|
|
|
|
# ─── TCP ─────────────────────────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class TCP:
|
|
src_port: int
|
|
dst_port: int
|
|
seq: int
|
|
ack: int
|
|
data_offset: int = 5 # 5 → 20-byte header
|
|
flags: int = 0
|
|
window: int = 0
|
|
checksum: int = 0
|
|
urg_ptr: int = 0
|
|
options: bytes = b''
|
|
payload: bytes = b''
|
|
|
|
@classmethod
|
|
def parse(cls, segment: bytes) -> 'TCP':
|
|
if len(segment) < 20:
|
|
raise ValueError(f'TCP too short: {len(segment)}')
|
|
src_port, dst_port, seq, ack = struct.unpack('!HHII', segment[:12])
|
|
off_flags = struct.unpack('!H', segment[12:14])[0]
|
|
data_offset = (off_flags >> 12) & 0xf
|
|
flags = off_flags & 0x1ff
|
|
window = struct.unpack('!H', segment[14:16])[0]
|
|
checksum = struct.unpack('!H', segment[16:18])[0]
|
|
urg_ptr = struct.unpack('!H', segment[18:20])[0]
|
|
header_len = data_offset * 4
|
|
if header_len < 20 or header_len > len(segment):
|
|
raise ValueError(f'bad TCP data_offset {data_offset}')
|
|
options = segment[20:header_len]
|
|
payload = segment[header_len:]
|
|
return cls(
|
|
src_port=src_port,
|
|
dst_port=dst_port,
|
|
seq=seq,
|
|
ack=ack,
|
|
data_offset=data_offset,
|
|
flags=flags,
|
|
window=window,
|
|
checksum=checksum,
|
|
urg_ptr=urg_ptr,
|
|
options=options,
|
|
payload=payload,
|
|
)
|
|
|
|
def to_bytes(self, src_ip: bytes, dst_ip: bytes) -> bytes:
|
|
# Pad options to 4-byte boundary.
|
|
opts = self.options
|
|
if len(opts) % 4:
|
|
opts = opts + b'\x00' * (4 - len(opts) % 4)
|
|
data_offset = (20 + len(opts)) // 4
|
|
off_flags = (data_offset << 12) | (self.flags & 0x1ff)
|
|
header = struct.pack(
|
|
'!HHIIHHHH',
|
|
self.src_port,
|
|
self.dst_port,
|
|
self.seq & 0xffffffff,
|
|
self.ack & 0xffffffff,
|
|
off_flags,
|
|
self.window,
|
|
0, # checksum placeholder
|
|
self.urg_ptr,
|
|
)
|
|
segment = header + opts + self.payload
|
|
cksum = tcp_udp_checksum(src_ip, dst_ip, IPPROTO_TCP, segment)
|
|
return segment[:16] + struct.pack('!H', cksum) + segment[18:]
|
|
|
|
|
|
def parse_tcp_options(opts: bytes) -> dict:
|
|
"""Parse TCP options into a {kind → value} dict. Handles MSS, NOP, EOL."""
|
|
out: dict = {}
|
|
i = 0
|
|
while i < len(opts):
|
|
kind = opts[i]
|
|
if kind == 0: # End of Option List
|
|
break
|
|
if kind == 1: # NOP
|
|
i += 1
|
|
continue
|
|
if i + 1 >= len(opts):
|
|
break
|
|
length = opts[i + 1]
|
|
if length < 2 or i + length > len(opts):
|
|
break
|
|
value = opts[i + 2:i + length]
|
|
if kind == 2 and length == 4: # MSS
|
|
out['mss'] = struct.unpack('!H', value)[0]
|
|
elif kind == 3 and length == 3: # Window Scale
|
|
out['wscale'] = value[0]
|
|
else:
|
|
out[kind] = value
|
|
i += length
|
|
return out
|
|
|
|
|
|
# ─── UDP ─────────────────────────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class UDP:
|
|
src_port: int
|
|
dst_port: int
|
|
length: int = 0
|
|
checksum: int = 0
|
|
payload: bytes = b''
|
|
|
|
@classmethod
|
|
def parse(cls, segment: bytes) -> 'UDP':
|
|
if len(segment) < 8:
|
|
raise ValueError(f'UDP too short: {len(segment)}')
|
|
src_port, dst_port, length, checksum = struct.unpack('!HHHH', segment[:8])
|
|
return cls(src_port, dst_port, length, checksum, segment[8:length])
|
|
|
|
def to_bytes(self, src_ip: bytes, dst_ip: bytes) -> bytes:
|
|
length = 8 + len(self.payload)
|
|
header = struct.pack('!HHHH', self.src_port, self.dst_port, length, 0)
|
|
segment = header + self.payload
|
|
cksum = tcp_udp_checksum(src_ip, dst_ip, IPPROTO_UDP, segment)
|
|
# UDP checksum 0 is "no checksum" — RFC 768 says when computed
|
|
# checksum is zero we send 0xffff to disambiguate.
|
|
if cksum == 0:
|
|
cksum = 0xffff
|
|
return segment[:6] + struct.pack('!H', cksum) + segment[8:]
|
|
|
|
|
|
# ─── ICMP ────────────────────────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class ICMP:
|
|
type: int = 0
|
|
code: int = 0
|
|
checksum: int = 0
|
|
rest: bytes = b'\x00\x00\x00\x00' # ID + sequence for echo
|
|
payload: bytes = b''
|
|
|
|
@classmethod
|
|
def parse(cls, segment: bytes) -> 'ICMP':
|
|
if len(segment) < 8:
|
|
raise ValueError(f'ICMP too short: {len(segment)}')
|
|
type_, code, checksum = struct.unpack('!BBH', segment[:4])
|
|
rest = segment[4:8]
|
|
payload = segment[8:]
|
|
return cls(type_, code, checksum, rest, payload)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
header_no_cksum = struct.pack('!BBH', self.type, self.code, 0) + self.rest + self.payload
|
|
cksum = internet_checksum(header_no_cksum)
|
|
return header_no_cksum[:2] + struct.pack('!H', cksum) + header_no_cksum[4:]
|
|
|
|
|
|
# ─── DHCP ────────────────────────────────────────────────────────────
|
|
#
|
|
# Layout per RFC 2131. Only fields we actually inspect/emit:
|
|
#
|
|
# op (1) htype (1) hlen (1) hops (1)
|
|
# xid (4)
|
|
# secs (2) flags (2)
|
|
# ciaddr (4) yiaddr (4) siaddr (4) giaddr (4)
|
|
# chaddr (16)
|
|
# sname (64) file (128)
|
|
# options (variable, prefixed by 4-byte magic cookie 0x63825363)
|
|
|
|
DHCP_MAGIC = b'\x63\x82\x53\x63'
|
|
|
|
|
|
@dataclass
|
|
class Dhcp:
|
|
op: int # 1=BOOTREQUEST, 2=BOOTREPLY
|
|
xid: int # transaction id
|
|
flags: int = 0
|
|
ciaddr: bytes = b'\x00\x00\x00\x00'
|
|
yiaddr: bytes = b'\x00\x00\x00\x00'
|
|
siaddr: bytes = b'\x00\x00\x00\x00'
|
|
giaddr: bytes = b'\x00\x00\x00\x00'
|
|
chaddr: bytes = b'\x00' * 16 # client hw addr (first 6 bytes are MAC)
|
|
options: dict = field(default_factory=dict) # {opt_code: value_bytes}
|
|
|
|
@classmethod
|
|
def parse(cls, payload: bytes) -> 'Dhcp':
|
|
if len(payload) < 240:
|
|
raise ValueError(f'DHCP too short: {len(payload)}')
|
|
op = payload[0]
|
|
xid = struct.unpack('!I', payload[4:8])[0]
|
|
flags = struct.unpack('!H', payload[10:12])[0]
|
|
ciaddr = payload[12:16]
|
|
yiaddr = payload[16:20]
|
|
siaddr = payload[20:24]
|
|
giaddr = payload[24:28]
|
|
chaddr = payload[28:44]
|
|
if payload[236:240] != DHCP_MAGIC:
|
|
raise ValueError('DHCP magic cookie missing')
|
|
opts: dict = {}
|
|
i = 240
|
|
while i < len(payload):
|
|
code = payload[i]
|
|
if code == 0: # pad
|
|
i += 1; continue
|
|
if code == 255: # end
|
|
break
|
|
if i + 1 >= len(payload):
|
|
break
|
|
length = payload[i + 1]
|
|
value = payload[i + 2:i + 2 + length]
|
|
opts[code] = value
|
|
i += 2 + length
|
|
return cls(op, xid, flags, ciaddr, yiaddr, siaddr, giaddr, chaddr, opts)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
head = struct.pack(
|
|
'!BBBBIHH4s4s4s4s16s64s128s',
|
|
self.op,
|
|
1, # htype = Ethernet
|
|
6, # hlen
|
|
0, # hops
|
|
self.xid,
|
|
0, # secs
|
|
self.flags,
|
|
self.ciaddr,
|
|
self.yiaddr,
|
|
self.siaddr,
|
|
self.giaddr,
|
|
self.chaddr,
|
|
b'', # sname
|
|
b'', # file
|
|
)
|
|
opts = bytearray(DHCP_MAGIC)
|
|
for code, value in self.options.items():
|
|
opts.append(code)
|
|
opts.append(len(value))
|
|
opts.extend(value)
|
|
opts.append(255) # end
|
|
return head + bytes(opts)
|
|
|
|
|
|
# ─── DNS (minimal — we only need to wrap a query/response) ───────────
|
|
|
|
@dataclass
|
|
class DnsMessage:
|
|
txid: int
|
|
flags: int
|
|
qd: list # list of (qname, qtype, qclass)
|
|
an: list = field(default_factory=list) # answers (name, type, class, ttl, rdata)
|
|
|
|
@classmethod
|
|
def parse(cls, payload: bytes) -> 'DnsMessage':
|
|
if len(payload) < 12:
|
|
raise ValueError(f'DNS too short: {len(payload)}')
|
|
txid, flags, qdcount, ancount, _nscount, _arcount = struct.unpack(
|
|
'!HHHHHH', payload[:12]
|
|
)
|
|
qd: list = []
|
|
offset = 12
|
|
for _ in range(qdcount):
|
|
qname, offset = _read_dns_name(payload, offset)
|
|
qtype, qclass = struct.unpack('!HH', payload[offset:offset + 4])
|
|
offset += 4
|
|
qd.append((qname, qtype, qclass))
|
|
an: list = []
|
|
for _ in range(ancount):
|
|
name, offset = _read_dns_name(payload, offset)
|
|
atype, aclass, ttl, rdlength = struct.unpack(
|
|
'!HHIH', payload[offset:offset + 10]
|
|
)
|
|
offset += 10
|
|
rdata = payload[offset:offset + rdlength]
|
|
offset += rdlength
|
|
an.append((name, atype, aclass, ttl, rdata))
|
|
return cls(txid=txid, flags=flags, qd=qd, an=an)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
head = struct.pack(
|
|
'!HHHHHH',
|
|
self.txid,
|
|
self.flags,
|
|
len(self.qd),
|
|
len(self.an),
|
|
0,
|
|
0,
|
|
)
|
|
body = bytearray()
|
|
for qname, qtype, qclass in self.qd:
|
|
body.extend(_write_dns_name(qname))
|
|
body.extend(struct.pack('!HH', qtype, qclass))
|
|
for name, atype, aclass, ttl, rdata in self.an:
|
|
body.extend(_write_dns_name(name))
|
|
body.extend(struct.pack('!HHIH', atype, aclass, ttl, len(rdata)))
|
|
body.extend(rdata)
|
|
return head + bytes(body)
|
|
|
|
|
|
def _read_dns_name(payload: bytes, offset: int) -> tuple[str, int]:
|
|
"""Parse a DNS name with simple compression-pointer support."""
|
|
labels: list[str] = []
|
|
seen_pointer = False
|
|
return_offset = offset
|
|
while offset < len(payload):
|
|
length = payload[offset]
|
|
if length == 0:
|
|
offset += 1
|
|
break
|
|
if length & 0xc0 == 0xc0:
|
|
# Pointer — top 2 bits set, next byte completes the offset.
|
|
if not seen_pointer:
|
|
return_offset = offset + 2
|
|
seen_pointer = True
|
|
offset = ((length & 0x3f) << 8) | payload[offset + 1]
|
|
continue
|
|
offset += 1
|
|
labels.append(payload[offset:offset + length].decode('ascii', errors='replace'))
|
|
offset += length
|
|
if not seen_pointer:
|
|
return_offset = offset
|
|
return ('.'.join(labels), return_offset)
|
|
|
|
|
|
def _write_dns_name(name: str) -> bytes:
|
|
out = bytearray()
|
|
if name:
|
|
for label in name.split('.'):
|
|
if not label:
|
|
continue
|
|
data = label.encode('ascii', errors='replace')[:63]
|
|
out.append(len(data))
|
|
out.extend(data)
|
|
out.append(0)
|
|
return bytes(out)
|
|
|
|
|
|
# ─── Convenience: build a complete L2/L3/L4 frame ────────────────────
|
|
|
|
def make_frame_ipv4(
|
|
dst_mac: bytes,
|
|
src_mac: bytes,
|
|
src_ip: bytes,
|
|
dst_ip: bytes,
|
|
protocol: int,
|
|
l4_payload: bytes,
|
|
ttl: int = 64,
|
|
ident: int = 0,
|
|
) -> bytes:
|
|
ipv4 = IPv4(
|
|
protocol=protocol,
|
|
src=src_ip,
|
|
dst=dst_ip,
|
|
ttl=ttl,
|
|
ident=ident,
|
|
payload=l4_payload,
|
|
)
|
|
return Ethernet(dst_mac, src_mac, ETHERTYPE_IPV4, ipv4.to_bytes()).to_bytes()
|
|
|
|
|
|
def make_frame_arp(
|
|
dst_mac: bytes,
|
|
src_mac: bytes,
|
|
arp: Arp,
|
|
) -> bytes:
|
|
return Ethernet(dst_mac, src_mac, ETHERTYPE_ARP, arp.to_bytes()).to_bytes()
|
|
|
|
|
|
__all__ = [
|
|
'Ethernet', 'Arp', 'IPv4', 'TCP', 'UDP', 'ICMP', 'Dhcp', 'DnsMessage',
|
|
'parse_tcp_options', 'make_frame_ipv4', 'make_frame_arp',
|
|
'DHCP_MAGIC', 'BROADCAST_MAC',
|
|
'IPPROTO_ICMP', 'IPPROTO_TCP', 'IPPROTO_UDP',
|
|
]
|