velxio/backend/app/services/picow_net/protocols.py

534 lines
17 KiB
Python
Raw Normal View History

"""
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',
]