Merge pull request #348 from dtrodrigues/aim

full test coverage for aim module, more descriptive tlv parsing error
kbandla · Mar 18, 2017 · 962cb22 · 962cb22
2 parents b0e08db + d38e0b5
commit 962cb22
Showing 1 changed file with 57 additions and 2 deletions.
diff --git a/dpkt/aim.py b/dpkt/aim.py
@@ -60,14 +60,69 @@ def tlv(buf):
     try:
         t, l = struct.unpack('>HH', buf[:n])
     except struct.error:
-        raise dpkt.UnpackError
+        raise dpkt.UnpackError('invalid type, length fields')
     v = buf[n:n + l]
     if len(v) < l:
-        raise dpkt.NeedData
+        raise dpkt.NeedData('%d left, %d needed' % (len(v), l))
     buf = buf[n + l:]
     return t, l, v, buf
 
 # TOC 1.0: http://jamwt.com/Py-TOC/PROTOCOL
 
 # TOC 2.0: http://www.firestuff.org/projects/firetalk/doc/toc2.txt
 
+def testAIM():
+    testdata = b'*\x02\xac\xf3\x00\x81\x00\x03\x00\x0b\x00\x00\xfaEUd\x0eusrnameremoved\x00\x00\x00\n\x00\x01\x00\x02\x12\x90\x00D\x00\x01\x00\x00\x03\x00\x04X\x90T6\x00E\x00\x04\x00\x00\x0f\x93\x00!\x00\x08\x00\x85\x00}\x00}\x00\x00\x00A\x00\x01\x00\x007\x00\x04\x00\x00\x00\x00\x00\r\x00\x00\x00\x19\x00\x00\x00\x1d\x00$\x00\x00\x00\x05\x02\x01\xd2\x04r\x00\x01\x00\x05\x02\x01\xd2\x04r\x00\x03\x00\x05+\x00\x00*\xcc\x00\x81\x00\x05+\x00\x00\x13\xf1'
+
+    flap = FLAP(testdata)
+    assert flap.ast == 0x2a
+    assert flap.type == 0x02
+    assert flap.seq == 44275
+    assert flap.len == 129
+    assert flap.data == b'\x00\x03\x00\x0b\x00\x00\xfaEUd\x0eusrnameremoved\x00\x00\x00\n\x00\x01\x00\x02\x12\x90\x00D\x00\x01\x00\x00\x03\x00\x04X\x90T6\x00E\x00\x04\x00\x00\x0f\x93\x00!\x00\x08\x00\x85\x00}\x00}\x00\x00\x00A\x00\x01\x00\x007\x00\x04\x00\x00\x00\x00\x00\r\x00\x00\x00\x19\x00\x00\x00\x1d\x00$\x00\x00\x00\x05\x02\x01\xd2\x04r\x00\x01\x00\x05\x02\x01\xd2\x04r\x00\x03\x00\x05+\x00\x00*\xcc\x00\x81\x00\x05+\x00\x00\x13\xf1'
+
+    snac = SNAC(flap.data)
+    assert snac.family == 3
+    assert snac.subtype == 11
+    assert snac.flags == 0
+    assert snac.reqid == 0xfa455564
+    assert snac.data == b'\x0eusrnameremoved\x00\x00\x00\n\x00\x01\x00\x02\x12\x90\x00D\x00\x01\x00\x00\x03\x00\x04X\x90T6\x00E\x00\x04\x00\x00\x0f\x93\x00!\x00\x08\x00\x85\x00}\x00}\x00\x00\x00A\x00\x01\x00\x007\x00\x04\x00\x00\x00\x00\x00\r\x00\x00\x00\x19\x00\x00\x00\x1d\x00$\x00\x00\x00\x05\x02\x01\xd2\x04r\x00\x01\x00\x05\x02\x01\xd2\x04r\x00\x03\x00\x05+\x00\x00*\xcc\x00\x81\x00\x05+\x00\x00\x13\xf1'
+
+    #skip over the buddyname and TLV count in Oncoming Buddy message
+    tlvdata = snac.data[19:]
+
+    tlvCount = 0
+    while tlvdata:
+        t, l, v, tlvdata = tlv(tlvdata)
+        tlvCount += 1
+        if tlvCount == 1:
+            # just check function return for first TLV
+            assert t == 0x01
+            assert l == 2
+            assert v == b'\x12\x90'
+            assert tlvdata == b'\x00D\x00\x01\x00\x00\x03\x00\x04X\x90T6\x00E\x00\x04\x00\x00\x0f\x93\x00!\x00\x08\x00\x85\x00}\x00}\x00\x00\x00A\x00\x01\x00\x007\x00\x04\x00\x00\x00\x00\x00\r\x00\x00\x00\x19\x00\x00\x00\x1d\x00$\x00\x00\x00\x05\x02\x01\xd2\x04r\x00\x01\x00\x05\x02\x01\xd2\x04r\x00\x03\x00\x05+\x00\x00*\xcc\x00\x81\x00\x05+\x00\x00\x13\xf1'
+
+    # make sure we extracted 10 TLVs
+    assert tlvCount == 10
+
+
+def testExceptions():
+    testdata = b'xxxxxx'
+    try:
+        flap = FLAP(testdata)
+    except dpkt.UnpackError as e:
+        assert str(e) == 'invalid FLAP header'
+    testdata = b'*\x02\x12\x34\x00\xff'
+    try:
+        flap = FLAP(testdata)
+    except dpkt.NeedData as e:
+        assert str(e) == '0 left, 255 needed'
+    try:
+        t, l, v, _ = tlv(b'x')
+    except dpkt.UnpackError as e:
+        assert str(e) == 'invalid type, length fields'
+
+    try:
+        t, l, v, _ = tlv(b'\x00\x01\x00\xff')
+    except dpkt.NeedData as e:
+        assert str(e) == '0 left, 255 needed'