BitStructures/src/examples/ip_codec.py at trunk · WibblyGhost/BitStructures · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
"""
Internet Protocol version 4 (TCP/IP protocol stack)
See RFC791

Transmission Control Protocol (TCP/IP protocol stack)
See RFC9293

User Datagram Protocol (TCP/IP protocol stack)
See RFC768
"""  # noqa: D400, D415

from math import ceil

from bitstructures import (
    Bits,
    Const,
    Enum,
    ExprAdapter,
    Flag,
    IpAddress,
    Optional,
    Padding,
    Struct,
    Switch,
)
from bitstructures.base.codec import Codec, GreedyBits

IPV4_HEADER = Struct(
    "version" / Const(Bits(4), const=4),
    "header_length"
    / ExprAdapter(
        # Indicates the length of the header in 32-bit words (minimum is 5, which equals 20 bytes).
        Bits(4),
        encoder=lambda obj: obj * 4,
        decoder=lambda obj: ceil(obj / 4),
    ),
    "tos"
    / Struct(
        "precedence" / Bits(3),
        "minimize_delay" / Flag(),
        "high_throuput" / Flag(),
        "high_reliability" / Flag(),
        "minimize_cost" / Flag(),
        Padding(1),
        embedded=False,
    ),
    "total_length" / Bits(16),
    "identification" / Bits(16),
    "flags" / Struct(Padding(1), "dont_fragment" / Flag(), "more_fragments" / Flag()),
    "fragment_offset" / Bits(13),
    "ttl" / Bits(8),
    "protocol" / Enum(8, ICMP=1, TCP=6, UDP=17),
    "checksum" / Bits(16),
    "source_ip" / IpAddress(Bits(32)),
    "destination_ip" / IpAddress(Bits(32)),
    "options" / Optional(Bits(lambda packet: packet.header_length - 20)),
)

TCP_HEADER = Struct(
    "source_port" / Bits(16),
    "destination_port" / Bits(16),
    "seq" / Bits(32),
    "ack" / Bits(32),
    "length" / ExprAdapter(Bits(4), encoder=lambda obj: obj * 4, decoder=lambda obj: ceil(obj / 4)),
    Padding(3),
    "flags"
    / Struct(
        "ns" / Flag(),
        "cwr" / Flag(),
        "ece" / Flag(),
        "urg" / Flag(),
        "ack" / Flag(),
        "psh" / Flag(),
        "rst" / Flag(),
        "syn" / Flag(),
        "fin" / Flag(),
        embedded=False,
    ),
    "window" / Bits(16),
    "checksum" / Bits(16),
    "urgent" / Bits(16),
    "options" / Optional(Bits(lambda packet: packet.length - 20)),
)

UDP_HEADER = Struct(
    "source_port" / Bits(16),
    "destination_port" / Bits(16),
    # Indicates the total length of the UDP header plus the payload.
    # The minimum value for this field is 8 (the header size), as there is always a header present.
    "length" / Bits(16),
    "checksum" / Bits(16),
)


# IP/UDP packet used for basic packet data tests, not in spec
IP_PACKET = Struct(
    IPV4_HEADER,
    "header"
    / Switch[str, Codec](lambda packet: packet.protocol, {"UDP": UDP_HEADER, "TCP": TCP_HEADER}),
    "payload" / GreedyBits(),
)


if __name__ == "__main__":
    from scapy.layers.inet import IP, UDP

    ip = IP(src="127.0.0.1", dst="8.8.8.8") / UDP(dport=53) / b"test_packet"
    ip_data = ip.build()
    print(f"IP: {ip_data!r}")
    stack = IP_PACKET.parse(ip_data)
    print(f"DECODED:\n{stack.pprint()}")
    container = stack.to_container()
    print(f"Container:\n{container!s}")
    raw = IP_PACKET.build(container)
    print(f"RAW:\n{raw!r}")
    print(f"{raw == ip_data=}")