mirror of
https://github.com/vivisect/vivisect
synced 2026-06-08 18:04:23 +00:00
James Gross
80de840881
A lot of cleanup things in prep for a python 3 transition. Getting rid of the old exception syntax, converting prints over to logging, cutting random scraps of code to be proper unit tests, cut away some older bits of code, etc. This still works in python2. It's just a lot of tidying up. There are no major functionality changes.
242 lines
7.0 KiB
Python
242 lines
7.0 KiB
Python
import binascii
|
|
import unittest
|
|
|
|
import vstruct
|
|
import vstruct.cparse as s_cparse
|
|
from vstruct.primitives import *
|
|
from vstruct.bitfield import *
|
|
|
|
from io import StringIO
|
|
|
|
|
|
class woot(vstruct.VStruct):
|
|
def __init__(self):
|
|
vstruct.VStruct.__init__(self)
|
|
self.lenfield = v_uint8()
|
|
self.strfield = v_str(size=0x20)
|
|
|
|
def pcb_lenfield(self):
|
|
self.vsGetField('strfield').vsSetLength(self.lenfield)
|
|
|
|
|
|
class AwesomeTest(s_cparse.CVStruct):
|
|
'''
|
|
struct awesome {
|
|
int x,z;
|
|
char stuff[20];
|
|
int y;
|
|
struct haha {
|
|
int blah;
|
|
} s;
|
|
int *q;
|
|
};
|
|
'''
|
|
|
|
|
|
class VStructTest(unittest.TestCase):
|
|
|
|
def test_autoparse(self):
|
|
awe = AwesomeTest()
|
|
awe.vsParse('XXXXZZZZhow cool is this?\x00\x00\x00YYYYblahQQQQ')
|
|
self.assertEqual(awe.x, 0x58585858)
|
|
self.assertEqual(awe.z, 0x5A5A5A5A)
|
|
self.assertEqual(awe.stuff, 'how cool is this?')
|
|
self.assertEqual(awe.y, 0x59595959)
|
|
self.assertEqual(awe.s.blah, 0x68616c62)
|
|
self.assertEqual(awe.q, 0x51515151)
|
|
|
|
def test_vstruct_basicstruct(self):
|
|
|
|
v = vstruct.VStruct()
|
|
v.uint8 = v_uint8(1)
|
|
v.uint16 = v_uint16(2)
|
|
v.uint24 = v_uint24(3)
|
|
v.uint32 = v_uint32(4)
|
|
v.uint64 = v_uint64(5)
|
|
v.vbytes = v_bytes(vbytes='ABCD')
|
|
|
|
answer = binascii.unhexlify('01020003000004000000050000000000000041424344')
|
|
self.assertEqual( v.vsEmit(), answer )
|
|
|
|
|
|
def test_vstruct_basicreasign(self):
|
|
v = vstruct.VStruct()
|
|
v.uint8 = v_uint8(1)
|
|
v.uint16 = v_uint16(2)
|
|
v.uint24 = v_uint24(3)
|
|
v.uint32 = v_uint32(4)
|
|
v.uint64 = v_uint64(5)
|
|
v.vbytes = v_bytes(vbytes='ABCD')
|
|
|
|
v.uint8 = 99
|
|
v.uint16 = 100
|
|
v.uint24 = 101
|
|
v.uint32 = 102
|
|
v.uint64 = 103
|
|
v.vbytes = '\x00\x00\x00\x00'
|
|
|
|
answer = binascii.unhexlify('63640065000066000000670000000000000000000000')
|
|
self.assertEqual( v.vsEmit(), answer )
|
|
|
|
|
|
def test_vstruct_fieldalign(self):
|
|
v = vstruct.VStruct()
|
|
v._vs_field_align = True
|
|
v.uint8 = v_uint8(0x42, bigend=True)
|
|
v.uint16 = v_uint16(0x4243, bigend=True)
|
|
v.uint24 = v_uint24(0x424344, bigend=True)
|
|
v.uint32 = v_uint32(0x42434445, bigend=True)
|
|
v.uint64 = v_uint64(0x4243444546474849, bigend=True)
|
|
|
|
answer = binascii.unhexlify('420042430000424344000000424344454243444546474849')
|
|
self.assertEqual( v.vsEmit(), answer )
|
|
|
|
def test_vstruct_fixedpartialasign(self):
|
|
v = vstruct.VStruct()
|
|
v.strfield = v_str(size=30)
|
|
v.unifield = v_wstr(size=30)
|
|
|
|
v.strfield = 'wootwoot!'
|
|
v.unifield = 'bazbaz'
|
|
|
|
answer = binascii.unhexlify('776f6f74776f6f7421000000000000000000000000000000000000000000620061007a00620061007a00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000')
|
|
|
|
self.assertEqual( v.vsEmit(), answer )
|
|
|
|
def test_vstruct_lengthcallback(self):
|
|
|
|
def updatelen(vs):
|
|
vs.vsGetField('strfield').vsSetLength(vs.lenfield)
|
|
|
|
v = vstruct.VStruct()
|
|
v.lenfield = v_uint8(0x30)
|
|
v.strfield = v_str(size=30)
|
|
v.vsAddParseCallback('lenfield', updatelen)
|
|
|
|
v.vsParse('\x01' + 'A' * 30)
|
|
self.assertEqual( v.vsEmit(), binascii.unhexlify('0141') )
|
|
|
|
|
|
def test_vstruct_classcallback(self):
|
|
v = woot()
|
|
v.vsParse('\x01' + 'A'*30)
|
|
self.assertEqual( v.vsEmit(), binascii.unhexlify('0141') )
|
|
|
|
def test_vstruct_parsefd(self):
|
|
v = woot()
|
|
sio = StringIO(('\x01' + 'A' * 30).decode('utf-8'))
|
|
v.vsParseFd(sio)
|
|
self.assertEqual( v.vsEmit(), binascii.unhexlify('0141') )
|
|
|
|
def test_vstruct_insertfield(self):
|
|
v = woot()
|
|
v.vsInsertField('ifield', v_uint8(), 'strfield')
|
|
v.vsParse('\x01BAAAAA')
|
|
self.assertEqual( v.vsEmit(), binascii.unhexlify('014241') )
|
|
|
|
def test_vstruct_floats(self):
|
|
|
|
v = vstruct.VStruct()
|
|
v.float4 = v_float()
|
|
v.float8 = v_double()
|
|
|
|
v.float4 = 99.3
|
|
v.float8 = -400.2
|
|
|
|
self.assertEqual( v.vsEmit(), binascii.unhexlify('9a99c64233333333330379c0') )
|
|
|
|
def test_vstruct_fastparse(self):
|
|
v = vstruct.VStruct()
|
|
v.x = v_uint8()
|
|
v.y = v_str(size=4)
|
|
v.z = v_uint32()
|
|
|
|
v.vsParse('BAAAAABCD', fast=True)
|
|
|
|
self.assertEqual( v.x, 0x42 )
|
|
self.assertEqual( v.y, 'AAAA' )
|
|
self.assertEqual( v.z, 0x44434241 )
|
|
|
|
def test_vstruct_fastparse_bigend(self):
|
|
v = vstruct.VStruct()
|
|
v.x = v_uint8()
|
|
v.y = v_str(size=4)
|
|
v.z = v_uint32(bigend=True)
|
|
|
|
v.vsParse('BAAAAABCD', fast=True)
|
|
|
|
self.assertEqual( v.x, 0x42 )
|
|
self.assertEqual( v.y, 'AAAA' )
|
|
self.assertEqual( v.z, 0x41424344 )
|
|
|
|
def test_vstruct_varray(self):
|
|
v = vstruct.VArray( [ v_uint8(i) for i in range(20) ] )
|
|
self.assertEqual( v[2], 2 )
|
|
v.vsParse('A' * 20)
|
|
self.assertEqual( v[2], 0x41 )
|
|
|
|
def test_bitfield(self):
|
|
v = VBitField()
|
|
v.vsAddField('w', v_bits(2))
|
|
v.vsAddField('x', v_bits(3))
|
|
v.vsAddField('y', v_bits(3))
|
|
v.vsAddField('z', v_bits(11))
|
|
v.vsAddField('a', v_bits(3))
|
|
|
|
v.vsAddField('stuff', v_bits(23))
|
|
v.vsAddField('pad', v_bits(3))
|
|
v.vsAddField('pad2', v_bits(6))
|
|
v.vsAddField('pad3', v_bits(2))
|
|
|
|
v.vsParse('AAAAAAA')
|
|
self.assertEqual(1, v.w)
|
|
self.assertEqual(0, v.x)
|
|
self.assertEqual(1, v.y)
|
|
self.assertEqual(0x20a, v.z)
|
|
self.assertEqual(0, v.a)
|
|
self.assertEqual(0x282828, v.stuff)
|
|
self.assertEqual(1, v.pad)
|
|
self.assertEqual(16, v.pad2)
|
|
self.assertEqual(1, v.pad3)
|
|
|
|
self.assertEqual('AAAAAAA', v.vsEmit())
|
|
|
|
v.vsParse('ABCDEFG')
|
|
self.assertEqual(1, v.w)
|
|
self.assertEqual(0, v.x)
|
|
self.assertEqual(1, v.y)
|
|
self.assertEqual(0x212, v.z)
|
|
self.assertEqual(0, v.a)
|
|
self.assertEqual(0x6888a8, v.stuff)
|
|
self.assertEqual(6, v.pad)
|
|
self.assertEqual(17, v.pad2)
|
|
self.assertEqual(3, v.pad3)
|
|
|
|
self.assertEqual('ABCDEFG', v.vsEmit())
|
|
|
|
v.vsParse('zxcvbnm')
|
|
self.assertEqual(1, v.w)
|
|
self.assertEqual(7, v.x)
|
|
self.assertEqual(2, v.y)
|
|
self.assertEqual(0x3c3, v.z)
|
|
self.assertEqual(0, v.a)
|
|
self.assertEqual(0x6ecc4d, v.stuff)
|
|
self.assertEqual(6, v.pad)
|
|
self.assertEqual(0x1b, v.pad2)
|
|
self.assertEqual(1, v.pad3)
|
|
|
|
self.assertEqual('zxcvbnm', v.vsEmit())
|
|
|
|
v.vsParse('asdfghj')
|
|
self.assertEqual(1, v.w)
|
|
self.assertEqual(4, v.x)
|
|
self.assertEqual(1, v.y)
|
|
self.assertEqual(0x39b, v.z)
|
|
self.assertEqual(1, v.a)
|
|
self.assertEqual(0xccced, v.stuff)
|
|
self.assertEqual(0, v.pad)
|
|
self.assertEqual(0x1a, v.pad2)
|
|
self.assertEqual(2, v.pad3)
|
|
|
|
self.assertEqual('asdfghj', v.vsEmit())
|