Test Failed
Pull Request — master (#4)
by macartur
05:24
created

IPv4.pack()   B

Complexity

Conditions 2

Size

Total Lines 24

Duplication

Lines 0
Ratio 0 %

Importance

Changes 3
Bugs 0 Features 0
Metric Value
dl 0
loc 24
c 3
b 0
f 0
rs 8.9713
cc 2
1
"""Basic Network packet types.
2
3
Defines and Implements Basic Network packet types , such as Ethertnet and LLDP.
4
"""
5
6
from pyof.foundation.base import GenericStruct
7
from pyof.foundation.basic_types import (
8
    BinaryData, HWAddress, IPAddress, UBInt8, UBInt16)
9
from pyof.foundation.constants import VLAN_TPID
10
from pyof.foundation.exceptions import PackException, UnpackException
11
12
__all__ = ('ARP', 'Ethernet', 'GenericTLV', 'IPv4', 'VLAN', 'TLVWithSubType',
13
           'LLDP')
14
15
16
class ARP(GenericStruct):
17
    """ARP packet "struct".
18
19
    Contains fields for an ARP packet's header and data.
20
    Designed for Ethernet and IPv4 only: needs to have some attributes changed
21
    for other HTYPE and PTYPE implementations.
22
    Must be encapsulated inside an Ethernet frame.
23
    """
24
25
    htype = UBInt16()
26
    ptype = UBInt16()
27
    hlen = UBInt8()
28
    plen = UBInt8()
29
    oper = UBInt16()
30
    sha = HWAddress()
31
    spa = IPAddress()
32
    tha = HWAddress()
33
    tpa = IPAddress()
34
35
    def __init__(self, htype=1, ptype=0x800, hlen=6, plen=4, oper=1,
36
                 sha='00:00:00:00:00:00', spa='0.0.0.0',
37
                 tha="00:00:00:00:00:00", tpa='0.0.0.0'):
38
        """Create a ARP with the parameters below.
39
40
        Args:
41
            htype (int): Hardware protocol type. Defaults to 1 for Ethernet.
42
            ptype (int): Network protocol type. Defaults to 0x800 for IPv4.
43
            hlen (int): Length of the hardware address. Defaults to 6 for MAC
44
                        addresses.
45
            plen (int): Length of the networking protocol address. Defaults to
46
                        4 for IPv4 addresses.
47
            oper (int): Determines the operation for this ARP packet. Must be 1
48
                        for ARP request or 2 for ARP reply. Defaults to 1.
49
            sha (str): Sender hardware address. Defaults to
50
                       '00:00:00:00:00:00'.
51
            spa (str): Sender protocol address. Defaults to '0.0.0.0'.
52
            tha (str): Target hardware address. Defaults to
53
                       '00:00:00:00:00:00'.
54
            tpa (str): Target protocol address. Defaults to '0.0.0.0'.
55
        """
56
        super().__init__()
57
        self.htype = htype
58
        self.ptype = ptype
59
        self.hlen = hlen
60
        self.plen = plen
61
        self.oper = oper
62
        self.sha = sha
63
        self.spa = spa
64
        self.tha = tha
65
        self.tpa = tpa
66
67
    def is_valid(self):
68
        """Assure the ARP contains Ethernet and IPv4 information."""
69
        return self.htype == 1 and self.ptype == 0x800
70
71
    def unpack(self, buff, offset=0):
72
        """Unpack a binary struct into this object's attributes.
73
74
        Return the values instead of the lib's basic types.
75
        Check if the protocols involved are Ethernet and IPv4. Other protocols
76
        are currently not supported.
77
78
        Args:
79
            buff (bytes): Binary buffer.
80
            offset (int): Where to begin unpacking.
81
82
        Raises:
83
            :exc:`~.exceptions.UnpackException`: If unpack fails.
84
85
        """
86
        super().unpack(buff, offset)
87
        if not self.is_valid():
88
            raise UnpackException("Unsupported protocols in ARP packet")
89
90
91
class VLAN(GenericStruct):
92
    """802.1q VLAN header."""
93
94
    #: tpid (:class:`UBInt16`): Tag Protocol Identifier
95
    tpid = UBInt16(VLAN_TPID)
96
    #: _tci (:class:`UBInt16`): Tag Control Information - has the
97
    #: Priority Code Point, DEI/CFI bit and the VLAN ID
98
    _tci = UBInt16()
99
100
    def __init__(self, pcp=None, cfi=None, vid=None):
101
        """Create a VLAN with the parameters below.
102
103
        If no arguments are set for a particular instance, it is interpreted as
104
        abscence of VLAN information, and the pack() method will return an
105
        empty binary string.
106
107
        Args:
108
            tpid (int): Tag Protocol Identifier. Defaults to 0x8100 for 802.1q.
109
            pcp (int): 802.1p Priority Code Point. Defaults to 0 for Best
110
                       Effort Queue.
111
            cfi (int): Canonical Format Indicator. Defaults to 0 for Ethernet.
112
            vid (int): VLAN ID. If no VLAN is specified, value is 0.
113
        """
114
        super().__init__()
115
        self.tpid = VLAN_TPID
116
        self.pcp = pcp
117
        self.cfi = cfi
118
        self.vid = vid
119
120
    def pack(self, value=None):
121
        """Pack the struct in a binary representation.
122
123
        Merge some fields to ensure correct packing.
124
125
        If no arguments are set for a particular instance, it is interpreted as
126
        abscence of VLAN information, and the pack() method will return an
127
        empty binary string.
128
129
        Returns:
130
            bytes: Binary representation of this instance.
131
132
        """
133
        if isinstance(value, type(self)):
134
            return value.pack()
135
136
        if self.pcp is None and self.cfi is None and self.vid is None:
137
            return b''
138
        self.pcp = self.pcp if self.pcp is not None else 0
139
        self.cfi = self.cfi if self.cfi is not None else 0
140
        self.vid = self.vid if self.vid is not None else 0
141
        self._tci = self.pcp << 13 | self.cfi << 12 | self.vid
142
        return super().pack()
143
144
    def _validate(self):
145
        """Assure this is a 802.1q VLAN header instance."""
146
        if self.tpid.value != VLAN_TPID:
147
            raise UnpackException
148
        return
149
150
    def unpack(self, buff, offset=0):
151
        """Unpack a binary struct into this object's attributes.
152
153
        Return the values instead of the lib's basic types.
154
155
        After unpacking, the abscence of a `tpid` value causes the assignment
156
        of None to the field values to indicate that there is no VLAN
157
        information.
158
159
        Args:
160
            buff (bytes): Binary buffer.
161
            offset (int): Where to begin unpacking.
162
163
        Raises:
164
            :exc:`~.exceptions.UnpackException`: If unpack fails.
165
166
        """
167
        super().unpack(buff, offset)
168
        if self.tpid.value:
169
            self._validate()
170
            self.tpid = self.tpid.value
171
            self.pcp = self._tci.value >> 13
172
            self.cfi = (self._tci.value >> 12) & 1
173
            self.vid = self._tci.value & 4095
174
        else:
175
            self.tpid = VLAN_TPID
176
            self.pcp = None
177
            self.cfi = None
178
            self.vid = None
179
180
181
class Ethernet(GenericStruct):
182
    """Ethernet "struct".
183
184
    Objects of this class represents an ethernet packet. It contains the
185
    'Ethernet header', composed by destination (MAC), source (MAC), type
186
    (EtherType)[1] and the payload of the packet, as binary data.
187
188
    This class does not consider the Ethernet 'Preamble' or the 'CRC'.
189
190
    There is also a get_hash method, that hashes the binary representation of
191
    the object so we can have a unique representation of the ethernet packet,
192
    so we can keep a track of ethernet packets being flooded over the network.
193
194
    [1] EtherTypes:
195
    http://www.iana.org/assignments/ieee-802-numbers/ieee-802-numbers.xhtml#ieee-802-numbers-1
196
    """
197
198
    destination = HWAddress()
199
    source = HWAddress()
200
    vlan = VLAN()
201
    ether_type = UBInt16()
202
    data = BinaryData()
203
204
    def __init__(self, destination=None,
205
                 source=None, vlan=VLAN(), ether_type=None, data=b''):
206
        """Create an instance and set its attributes.
207
208
        Args:
209
            destination (:class:`~pyof.foundation.basic_types.HWAddress`):
210
                The final destination MAC address.
211
            source (:class:`~pyof.foundation.basic_types.HWAddress`):
212
                The source Mac address of the packet.
213
            ether_type (:class:`~pyof.foundation.basic_types.UBInt16`):
214
                The EtherType of packet.
215
            data (:class:`~pyof.foundation.basic_types.BinaryData`):
216
                The content of the packet in binary format.
217
        """
218
        super().__init__()
219
        self.destination = destination
220
        self.source = source
221
        self.vlan = vlan
222
        self.ether_type = ether_type
223
        self.data = data
224
225
    def get_hash(self):
226
        """Calculate a hash and returns it.
227
228
        Returns:
229
            int: Integer value that identifies this instance.
230
231
        """
232
        return hash(self.pack())
233
234
    def unpack(self, buff, offset=0):
235
        """Unpack a binary message into this object's attributes.
236
237
        Unpack the binary value *buff* and update this object attributes based
238
        on the results.
239
240
        Ethernet headers may have VLAN tags. If no VLAN tag is found, a
241
        'wildcard VLAN tag' is inserted to assure correct unpacking.
242
243
        Args:
244
            buff (bytes): Binary data package to be unpacked.
245
            offset (int): Where to begin unpacking.
246
247
        Raises:
248
            UnpackException: If there is a struct unpacking error.
249
250
        """
251
        # Checking if the EtherType bytes are actually equal to VLAN_TPID -
252
        # indicating that the packet is tagged. If it is not, we insert the
253
        # equivalent to 'NULL VLAN data' (\x00\x00\x00\x00) to enable the
254
        # correct unpacking process.
255
        if buff[12:16] != VLAN_TPID.to_bytes(2, 'big'):
256
            buff = buff[0:12] + b'\x00\x00\x00\x00' + buff[12:]
257
258
        super().unpack(buff, offset)
259
260
261
class GenericTLV(GenericStruct):
262
    """TLV structure of LLDP packets.
263
264
    This is a Type, Length and Value (TLV) struct.
265
266
    The LLDP/TLV definition states that the Type field have 7 bits, while
267
    the length have 9 bits. The Value must be between 0-511 octets.
268
269
    Internally, on the instances of this class, the Type is a integer
270
    (0-127) and the Length is dynamically calculated based on the current
271
    type and value.
272
    """
273
274
    def __init__(self, tlv_type=127, value=BinaryData()):
275
        """Create an instance and set its attributes.
276
277
        Args:
278
            tlv_type (int): Type used by this class. Defaults to 127.
279
            value (:class:`~pyof.foundation.basic_types.BinaryData`):
280
                Value stored by GenericTLV.
281
        """
282
        super().__init__()
283
        self.tlv_type = tlv_type
284
        self._value = value
285
286
    @property
287
    def value(self):
288
        """Return the value stored by GenericTLV.
289
290
        Returns:
291
            :class:`~pyof.foundation.basic_types.BinaryData`:
292
                Value stored by GenericTLV.
293
294
        """
295
        return self._value
296
297
    @property
298
    def length(self):
299
        """Return the length of value stored by GenericTLV.
300
301
        Returns:
302
            int: Value length in bytes.
303
304
        """
305
        return len(self.value.pack())
306
307
    @property
308
    def header(self):
309
        """Header of the TLV Packet.
310
311
        The header is composed by the Type (7 bits) and Length (9 bits),
312
        summing up 16 bits. To achieve that, we need to do some bitshift
313
        operations.
314
315
        Returns:
316
            :class:`~pyof.foundation.basic_types.UBInt16`:
317
                Result after all operations.
318
319
        """
320
        return UBInt16(((self.tlv_type & 127) << 9) | (self.length & 511))
321
322
    def pack(self, value=None):
323
        """Pack the TLV in a binary representation.
324
325
        Returns:
326
            bytes: Binary representation of the struct object.
327
328
        Raises:
329
            :exc:`~.exceptions.ValidationError`: If validation fails.
330
331
        """
332
        if value is None:
333
            output = self.header.pack()
334
            output += self.value.pack()
335
            return output
336
337
        elif isinstance(value, type(self)):
338
            return value.pack()
339
        else:
340
            msg = "{} is not an instance of {}".format(value,
341
                                                       type(self).__name__)
342
            raise PackException(msg)
343
344
    def unpack(self, buff, offset=0):
345
        """Unpack a binary message into this object's attributes.
346
347
        Unpack the binary value *buff* and update this object attributes based
348
        on the results.
349
350
        Args:
351
            buff (bytes): Binary data package to be unpacked.
352
            offset (int): Where to begin unpacking.
353
354
        Raises:
355
            Exception: If there is a struct unpacking error.
356
357
        """
358
        header = UBInt16()
359
        header.unpack(buff[offset:offset+2])
360
        self.tlv_type = header.value >> 9
361
        length = header.value & 511
362
        begin, end = offset + 2, offset + 2 + length
363
        self._value = BinaryData(buff[begin:end])
364
365
    def get_size(self, value=None):
366
        """Return struct size.
367
368
        Returns:
369
            int: Returns the struct size based on inner attributes.
370
371
        """
372
        if isinstance(value, type(self)):
373
            return value.get_size()
374
375
        return 2 + self.length
376
377
378
class IPv4(GenericStruct):
379
    """IPv4 packet "struct".
380
381
    Contains all fields of an IP version 4 packet header, plus the upper layer
382
    content as binary data.
383
    Some of the fields were merged together because of their size being
384
    inferior to 8 bits. They are represented as a single class attribute, but
385
    pack/unpack methods will take into account the values in individual
386
    instance attributes.
387
    """
388
389
    #: _version_ihl (:class:`UBInt8`): IP protocol version + Internet Header
390
    #: Length (words)
391
    _version_ihl = UBInt8()
392
    #: _dscp_ecn (:class:`UBInt8`): Differentiated Services Code Point
393
    #: (ToS - Type of Service) + Explicit Congestion Notification
394
    _dscp_ecn = UBInt8()
395
    #: length (:class:`UBInt16`): IP packet length (bytes)
396
    length = UBInt16()
397
    #: identification (:class:`UBInt16`): Packet ID - common to all fragments
398
    identification = UBInt16()
399
    #: _flags_offset (:class:`UBInt16`): Fragmentation flags + fragmentation
400
    #: offset
401
    _flags_offset = UBInt16()
402
    #: ttl (:class:`UBInt8`): Packet time-to-live
403
    ttl = UBInt8()
404
    #: protocol (:class:`UBInt8`): Upper layer protocol number
405
    protocol = UBInt8()
406
    #: checksum (:class:`UBInt16`): Header checksum
407
    checksum = UBInt16()
408
    #: source (:class:`IPAddress`): Source IPv4 address
409
    source = IPAddress()
410
    #: destination (:class:`IPAddress`): Destination IPv4 address
411
    destination = IPAddress()
412
    #: options (:class:`BinaryData`): IP Options - up to 320 bits, always
413
    #: padded to 32 bits
414
    options = BinaryData()
415
    #: data (:class:`BinaryData`): Packet data
416
    data = BinaryData()
417
418
    def __init__(self, version=4, ihl=5, dscp=0, ecn=0, length=0,
419
                 identification=0, flags=0, offset=0, ttl=255, protocol=0,
420
                 checksum=0, source="0.0.0.0", destination="0.0.0.0",
421
                 options=b'', data=b''):
422
        """Instanciate a IPv4 with the parameters below.
423
424
        Args:
425
            version (int): IP protocol version. Defaults to 4.
426
            ihl (int): Internet Header Length. Default is 5.
427
            dscp (int): Differentiated Service Code Point. Defaults to 0.
428
            ecn (int): Explicit Congestion Notification. Defaults to 0.
429
            length (int): IP packet length in bytes. Defaults to 0.
430
            identification (int): Packet Id. Defaults to 0.
431
            flags (int): IPv4 Flags. Defults 0.
432
            offset (int): IPv4 offset. Defaults to 0.
433
            ttl (int): Packet time-to-live. Defaults to 255
434
            protocol (int): Upper layer protocol number. Defaults to 0.
435
            checksum (int): Header checksum. Defaults to 0.
436
            source (str): Source IPv4 address. Defaults to "0.0.0.0"
437
            destination (str): Destination IPv4 address. Defaults to "0.0.0.0"
438
            options (bytes): IP options. Defaults to empty bytes.
439
            data (bytes): Packet data. Defaults to empty bytes.
440
        """
441
        super().__init__()
442
        self.version = version
443
        self.ihl = ihl
444
        self.dscp = dscp
445
        self.ecn = ecn
446
        self.length = length
447
        self.identification = identification
448
        self.flags = flags
449
        self.offset = offset
450
        self.ttl = ttl
451
        self.protocol = protocol
452
        self.checksum = checksum
453
        self.source = source
454
        self.destination = destination
455
        self.options = options
456
        self.data = data
457
458
    def _update_checksum(self):
459
        """Update the packet checksum to enable integrity check."""
460
        source_list = [int(octet) for octet in self.source.split(".")]
461
        destination_list = [int(octet) for octet in
462
                            self.destination.split(".")]
463
        source_upper = (source_list[0] << 8) + source_list[1]
464
        source_lower = (source_list[2] << 8) + source_list[3]
465
        destination_upper = (destination_list[0] << 8) + destination_list[1]
466
        destination_lower = (destination_list[2] << 8) + destination_list[3]
467
468
        block_sum = ((self._version_ihl << 8 | self._dscp_ecn) + self.length +
469
                     self.identification + self._flags_offset +
470
                     (self.ttl << 8 | self.protocol) + source_upper +
471
                     source_lower + destination_upper + destination_lower)
472
473
        while block_sum > 65535:
474
            carry = block_sum >> 16
475
            block_sum = (block_sum & 65535) + carry
476
477
        self.checksum = ~block_sum & 65535
478
479
    def pack(self, value=None):
480
        """Pack the struct in a binary representation.
481
482
        Merge some fields to ensure correct packing.
483
484
        Returns:
485
            bytes: Binary representation of this instance.
486
487
        """
488
        # Set the correct IHL based on options size
489
        if self.options:
490
            self.ihl += int(len(self.options) / 4)
491
492
        # Set the correct packet length based on header length and data
493
        self.length = int(self.ihl * 4 + len(self.data))
494
495
        self._version_ihl = self.version << 4 | self.ihl
496
        self._dscp_ecn = self.dscp << 2 | self.ecn
497
        self._flags_offset = self.flags << 13 | self.offset
498
499
        # Set the checksum field before packing
500
        self._update_checksum()
501
502
        return super().pack()
503
504
    def unpack(self, buff, offset=0):
505
        """Unpack a binary struct into this object's attributes.
506
507
        Return the values instead of the lib's basic types.
508
509
        Args:
510
            buff (bytes): Binary buffer.
511
            offset (int): Where to begin unpacking.
512
513
        Raises:
514
            :exc:`~.exceptions.UnpackException`: If unpack fails.
515
516
        """
517
        super().unpack(buff, offset)
518
519
        self.version = self._version_ihl.value >> 4
520
        self.ihl = self._version_ihl.value & 15
521
        self.dscp = self._dscp_ecn.value >> 2
522
        self.ecn = self._dscp_ecn.value & 3
523
        self.length = self.length.value
524
        self.identification = self.identification.value
525
        self.flags = self._flags_offset.value >> 13
526
        self.offset = self._flags_offset.value & 8191
527
        self.ttl = self.ttl.value
528
        self.protocol = self.protocol.value
529
        self.checksum = self.checksum.value
530
        self.source = self.source.value
531
        self.destination = self.destination.value
532
533
        if self.ihl > 5:
534
            options_size = (self.ihl - 5) * 4
535
            self.data = self.options.value[options_size:]
536
            self.options = self.options.value[:options_size]
537
        else:
538
            self.data = self.options.value
539
            self.options = b''
540
541
542
class TLVWithSubType(GenericTLV):
543
    """Modify the :class:`GenericTLV` to a Organization Specific TLV structure.
544
545
    Beyond the standard TLV (type, length, value), we can also have a more
546
    specific structure, with the :attr:`value` field being splitted into a
547
    :attr:`sub_type` field and a new :attr:`sub_value` field.
548
    """
549
550
    def __init__(self, tlv_type=1, sub_type=7, sub_value=BinaryData()):
551
        """Create an instance and set its attributes.
552
553
        Args:
554
            tlv_type (int): Type used by this class. Defaults to 1.
555
            sub_type (int): Sub type value used by this class. Defaults to 7.
556
            sub_value (:class:`~pyof.foundation.basic_types.BinaryData`):
557
                Data stored by TLVWithSubType. Defaults to empty BinaryData.
558
        """
559
        super().__init__(tlv_type)
560
        self.sub_type = sub_type
561
        self.sub_value = sub_value
562
563
    @property
564
    def value(self):
565
        """Return sub type and sub value as binary data.
566
567
        Returns:
568
            :class:`~pyof.foundation.basic_types.BinaryData`:
569
                BinaryData calculated.
570
571
        """
572
        binary = UBInt8(self.sub_type).pack() + self.sub_value.pack()
573
        return BinaryData(binary)
574
575
    def unpack(self, buff, offset=0):
576
        """Unpack a binary message into this object's attributes.
577
578
        Unpack the binary value *buff* and update this object attributes based
579
        on the results.
580
581
        Args:
582
            buff (bytes): Binary data package to be unpacked.
583
            offset (int): Where to begin unpacking.
584
585
        Raises:
586
            Exception: If there is a struct unpacking error.
587
588
        """
589
        header = UBInt16()
590
        header.unpack(buff[offset:offset+2])
591
        self.tlv_type = header.value >> 9
592
        length = header.value & 511
593
        begin, end = offset + 2, offset + 2 + length
594
        sub_type = UBInt8()
595
        sub_type.unpack(buff[begin:begin+1])
596
        self.sub_type = sub_type.value
597
        self.sub_value = BinaryData(buff[begin+1:end])
598
599
600
class LLDP(GenericStruct):
601
    """LLDP class.
602
603
    Build a LLDP packet with TLVSubtype and Generic Subtypes.
604
605
    It contains a chassis_id TLV, a port_id TLV, a TTL (Time to live) and
606
    another TLV to represent the end of the LLDP Packet.
607
    """
608
609
    #: chassis_id (:class:`~TLVWithSubType`) with tlv_type = 1 and sub_type = 7
610
    chassis_id = TLVWithSubType(tlv_type=1, sub_type=7)
611
    #: port_id (:class:`TLVWithSubType`) with tlv = 2 and sub_type = 7
612
    port_id = TLVWithSubType(tlv_type=2, sub_type=7)
613
    #: TTL (:class:`GenericTLV`) time is given in seconds, between 0 and 65535,
614
    #: with tlv_type = 3
615
    ttl = GenericTLV(tlv_type=3, value=UBInt16(120))
616
    # We are not using list of tlvs for now
617
    # tlvs = ListOfTLVs()
618
    #: end (:class:`GenericTLV`) with tlv_type = 0
619
    end = GenericTLV(tlv_type=0)
620