CS305-2024Spring-FinalProject/Header.py at main · SUSTech-HPCLab/CS305-2024Spring-FinalProject · 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
class RDTHeader():
    def __init__(self, SYN: int = 0, FIN: int = 0, ACK: int = 0, SEQ_num: int = 0, ACK_num: int = 0, LEN: int = 0, CHECKSUM: int = 0, PAYLOAD = None, RWND: int = 0)  -> None:
        self.test_case = 0                      # Indicate the test case that will be used

        self.SYN = SYN                          # 1 bytes
        self.FIN = FIN                          # 1 bytes
        self.ACK = ACK                          # 1 bytes
        self.SEQ_num = SEQ_num                  # 4 bytes
        self.ACK_num = ACK_num                  # 4 bytes
        self.LEN = LEN                          # 4 bytes
        self.CHECKSUM = CHECKSUM                # 2 bytes
        self.PAYLOAD = PAYLOAD                  # Data LEN bytes
        # self.CWND = CWND                      # Congestion window size 4 bytes
        self.RWND = RWND                        # Notification window size 4 bytes
        self.Reserved = 0                       # Reserved field for any attribte you need.

        self.Source_address = [127,0,0,1,12334] # Souce ip and port
        self.Target_address = [127,0,0,1,12345] # Target ip and port


    def to_bytes(self):
        test_case = self.test_case.to_bytes(1, 'big')
        Source_address = self.Source_address[0].to_bytes(1, 'big') + self.Source_address[1].to_bytes(1, 'big') + \
                            self.Source_address[2].to_bytes(1, 'big') + self.Source_address[3].to_bytes(1, 'big') + \
                            self.Source_address[4].to_bytes(2, 'big')

        Target_address = self.Target_address[0].to_bytes(1, 'big') + self.Target_address[1].to_bytes(1, 'big') + \
                            self.Target_address[2].to_bytes(1, 'big') + self.Target_address[3].to_bytes(1, 'big') + \
                            self.Target_address[4].to_bytes(2, 'big')

        SYN = self.SYN.to_bytes(1, 'big')
        FIN = self.FIN.to_bytes(1, 'big')
        ACK = self.ACK.to_bytes(1, 'big')
        SEQ_num = self.SEQ_num.to_bytes(4, 'big')
        ACK_num = self.ACK_num.to_bytes(4, 'big')
        LEN = self.LEN.to_bytes(4, 'big')
        RWND = self.RWND.to_bytes(4, 'big')
        CHECKSUM = self.CHECKSUM.to_bytes(2, 'big')
        PAYLOAD = self.PAYLOAD.encode() if isinstance(self.PAYLOAD, str) else "".encode()
        Reserved = self.Reserved.to_bytes(8, 'big')

        return b''.join([test_case, Source_address, Target_address, SYN, FIN, ACK, SEQ_num, ACK_num, LEN, RWND, CHECKSUM, Reserved, PAYLOAD])


    def from_bytes(self, data):
        self.test_case = data[0]
        Source_address = []
        for i in range(4):
            Source_address.append(data[i + 1])
        Source_address.append(int.from_bytes(data[5:7], 'big'))
        self.Source_address = Source_address

        Target_address = []
        for i in range(4):
            Target_address.append(data[i + 7])
        Target_address.append(int.from_bytes(data[11:13], 'big'))
        self.Target_address = Target_address


        self.SYN = data[13]
        self.FIN = data[14]
        self.ACK = data[15]
        self.SEQ_num = int.from_bytes(data[16:20], 'big')
        self.ACK_num = int.from_bytes(data[20:24], 'big')
        self.LEN = int.from_bytes(data[24:28], 'big')
        self.RWND = int.from_bytes(data[28:32], 'big')
        self.CHECKSUM = int.from_bytes(data[32:34], 'big')
        self.Reserved = int.from_bytes(data[34:42], 'big')

        self.PAYLOAD = data[42:].decode()

        return self

    def assign_address(self, Source_address: tuple, Target_address: tuple):
        source_ip = Source_address[0].split('.')
        target_ip = Target_address[0].split('.')
        for i in range(4):
            self.Source_address[i] = int(source_ip[i])
            self.Target_address[i] = int(target_ip[i])
        self.Source_address[4] = Source_address[1]
        self.Target_address[4] = Target_address[1]

    @property
    def src(self):
        return (f"{self.Source_address[0]}.{self.Source_address[1]}.{self.Source_address[2]}.{self.Source_address[3]}", self.Source_address[4])

    @property
    def tgt(self):
        return (f"{self.Target_address[0]}.{self.Target_address[1]}.{self.Target_address[2]}.{self.Target_address[3]}", self.Target_address[4])