tcp、udp报文的python代码实现

import struct

import socket
class UDP:
    def __init__(self, src_port, dst_port, payload):
        self.src_port = src_port
        self.dst_port = dst_port
        self.payload = payload

    def pack(self, src_ip, dst_ip):
        # UDP头部长度为8字节
        length = 8 + len(self.payload)
        # UDP头部（源端口号，目的端口号，长度，校验和）
        udp_header = struct.pack('!HHHH', self.src_port, self.dst_port, length, 0)

        # 伪首部，用于计算UDP校验和
        pseudo_header = struct.pack('!4s4sBBH',
                                    bytes(map(int, src_ip.split('.'))),
                                    bytes(map(int, dst_ip.split('.'))),
                                    0,  # 暂时填充为0
                                    socket.IPPROTO_UDP,
                                    length)

        # UDP数据报（UDP头部 + 数据）
        packet = udp_header + self.payload.encode()

        # 计算UDP校验和
        if length % 2 == 1:  # 补充成偶数字节
            packet += b'\0'
        chksum = self.checksum(pseudo_header + packet)
        udp_header = struct.pack('!HHHH', self.src_port, self.dst_port, length, chksum)

        return udp_header + self.payload.encode()

    def checksum(self, data):
        # UDP校验和算法
        length = len(data)
        if length % 2 == 1:
            data += b'\0'
            length += 1

        s = 0
        for i in range(0, length, 2):
            w = (data[i] << 8) + (data[i + 1])
            s += w

        s = (s >> 16) + (s & 0xffff)
        s += (s >> 16)
        s = ~s & 0xffff

        return s
if __name__=="__main__":

    udp_packet = UDP(1234, 5688, "Hello UDP!").pack('192.168.198.128', '192.168.198.1')
    print(udp_packet)

上述是udp报文的生成，以下是tcp报文的生成。

import struct
import random


class TCP:
    def __init__(self, src_port, dst_port, seq_num, ack_num, flags, window_size, data=b''):
        self.src_port = src_port
        self.dst_port = dst_port
        self.seq_num = seq_num
        self.ack_num = ack_num
        self.data_offset = 5  # TCP header size / 4
        self.reserved = 0
        self.flags = flags
        self.window_size = window_size
        self.checksum = 0  # will be calculated later
        self.urgent_ptr = 0
        self.options = b''
        self.payload = data

    def pack(self, src_ip, dst_ip):
        # Pseudo header for checksum calculation
        pseudo_header = struct.pack('!4s4sBBH',
                                    bytes.fromhex(src_ip),
                                    bytes.fromhex(dst_ip),
                                    0,  # Reserved
                                    6,  # Protocol number for TCP
                                    self.length(),
                                    )

        # TCP header
        tcp_header = struct.pack('!HHIIBBHHH',
                                 self.src_port,
                                 self.dst_port,
                                 self.seq_num,
                                 self.ack_num,
                                 (self.data_offset << 4) | self.reserved,
                                 self.flags,
                                 self.window_size,
                                 self.checksum,
                                 self.urgent_ptr,
                                 )

        # Calculate checksum
        self.checksum = self.calc_checksum(pseudo_header + tcp_header + self.payload)

        # Pack TCP header with updated checksum
        tcp_header = struct.pack('!HHIIBBHHH',
                                 self.src_port,
                                 self.dst_port,
                                 self.seq_num,
                                 self.ack_num,
                                 (self.data_offset << 4) | self.reserved,
                                 self.flags,
                                 self.window_size,
                                 self.checksum,
                                 self.urgent_ptr,
                                 )

        return tcp_header + self.payload

    def length(self):
        return self.data_offset * 4 + len(self.payload)

    def calc_checksum(self, data):
        """
        Reference: https://tools.ietf.org/html/rfc793#page-15
        """
        # Pad data if the length is odd
        if len(data) % 2 == 1:
            data += b'\x00'

        # Calculate the checksum
        checksum = 0
        for i in range(0, len(data), 2):
            checksum += (data[i] << 8) | data[i+1]

        checksum = (checksum >> 16) + (checksum & 0xffff)
        checksum += (checksum >> 16)

        return ~checksum & 0xffff


# Example usage
if __name__ == '__main__':
    tcp = TCP(src_port=1234,
              dst_port=5679,
              seq_num=123456,
              ack_num=0,
              flags=2,  # SYN
              window_size=8192,
              data=b'\x04\xd2\x16/\x00\x01\xe2@\x00\x00\x00\x00P\x02 \x00\x85K\x00\x00')

    packet = tcp.pack(src_ip='0a0a0a0a', dst_ip='0b0b0b0b')
    packet_hex = packet.hex()

    print(f'TCP Packet: {packet_hex}')
    print(packet)
b'\x04\xd2\x16/\x00\x01\xe2@\x00\x00\x00\x00P\x02 \x00\x00\x00\x00\x00'