1# $Id: ah.py 34 2007-01-28 07:54:20Z dugsong $
2# -*- coding: utf-8 -*-
3
4"""Authentication Header."""
5from __future__ import absolute_import
6
7from . import dpkt
8from . import ip
9
10
11class AH(dpkt.Packet):
12    """Authentication Header.
13
14    TODO: Longer class information....
15
16    Attributes:
17        __hdr__: Header fields of AH.
18        auth: Authentication body.
19        data: Message data.
20    """
21
22    __hdr__ = (
23        ('nxt', 'B', 0),
24        ('len', 'B', 0),  # payload length
25        ('rsvd', 'H', 0),
26        ('spi', 'I', 0),
27        ('seq', 'I', 0)
28    )
29    auth = b''
30
31    def unpack(self, buf):
32        dpkt.Packet.unpack(self, buf)
33        self.auth = self.data[:self.len]
34        buf = self.data[self.len:]
35
36        try:
37            self.data = ip.IP.get_proto(self.nxt)(buf)
38            setattr(self, self.data.__class__.__name__.lower(), self.data)
39        except (KeyError, dpkt.UnpackError):
40            self.data = buf
41
42    def __len__(self):
43        return self.__hdr_len__ + len(self.auth) + len(self.data)
44
45    def __bytes__(self):
46        return self.pack_hdr() + bytes(self.auth) + bytes(self.data)
47
48
49def test_default_creation():
50    ah = AH()
51    assert ah.nxt == 0
52    assert ah.len == 0
53    assert ah.rsvd == 0
54    assert ah.spi == 0
55    assert ah.seq == 0
56    assert len(ah) == ah.__hdr_len__
57    assert bytes(ah) == b'\x00' * 12
58
59
60def test_creation_from_buf():
61    from binascii import unhexlify
62    buf_ip = unhexlify(
63        '04'  # IP
64        '0000000000000000000000'
65        '4500002200000000401172c001020304'
66        '01020304006f00de000ebf35666f6f626172'
67    )
68
69    ah = AH(buf_ip)
70    assert ah.nxt == 4  # IP
71    assert isinstance(ah.data, ip.IP)
72    assert len(ah) == 46
73    assert bytes(ah) == buf_ip
74
75    buf_not_ip = unhexlify(
76        '37'  # Not registered
77        '0000000000000000000000'
78        '4500002200000000401172c001020304'
79        '01020304006f00de000ebf35666f6f626172'
80    )
81    ah_not_ip = AH(buf_not_ip)
82    assert ah_not_ip.nxt == 0x37
83    assert isinstance(ah_not_ip.data, bytes)
84    assert len(ah_not_ip) == 46
85    assert bytes(ah_not_ip) == buf_not_ip
86