mirror of
https://github.com/vivisect/vivisect
synced 2026-06-08 18:04:23 +00:00
7a883f4bbf
* Migrate away from deprecated `assertRaisesRegexp` In Python 3.12 `assertRaisesRegexp` will no longer work. Replace it with `assertRaisesRegex`. * Migrate deprecated threading.currentThread * Migrate from deprecated `threading.Thread.setDaemon`
144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
import unittest
|
|
|
|
import vstruct
|
|
import vstruct.primitives as p
|
|
|
|
class NNestedStruct(vstruct.VStruct):
|
|
def __init__(self):
|
|
vstruct.VStruct.__init__(self)
|
|
|
|
self.alpha = p.v_uint64()
|
|
self.beta = vstruct.VStruct()
|
|
|
|
class NestedStruct(vstruct.VStruct):
|
|
def __init__(self):
|
|
vstruct.VStruct.__init__(self)
|
|
|
|
self.foo = p.v_bytes(size=3)
|
|
self.bar = p.v_uint32()
|
|
self.baz = p.v_bytes(size=256)
|
|
self.faz = NNestedStruct()
|
|
|
|
class TestStruct(vstruct.VStruct):
|
|
def __init__(self):
|
|
vstruct.VStruct.__init__(self)
|
|
|
|
self.one = p.v_uint32()
|
|
self.two = NestedStruct()
|
|
self.three = p.v_uint32()
|
|
self.four = p.v_bytes(size=100)
|
|
|
|
class VStructTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.s = TestStruct()
|
|
|
|
def test_vsGetFieldByOffset_start(self):
|
|
fname, fobj = self.s.vsGetFieldByOffset(0)
|
|
self.assertEqual(fname, 'one')
|
|
self.assertIs(fobj, self.s.vsGetField('one'))
|
|
|
|
def test_vsGetFieldByOffset_end(self):
|
|
fname, fobj = self.s.vsGetFieldByOffset(len(self.s) - 1)
|
|
self.assertEqual(fname, 'four')
|
|
self.assertIs(fobj, self.s.vsGetField('four'))
|
|
|
|
def test_vsGetFieldByOffset_maxval_neg(self):
|
|
with self.assertRaisesRegex(Exception, 'Invalid Offset Specified'):
|
|
tup = self.s.vsGetFieldByOffset(-1)
|
|
|
|
def test_vsGetFieldByOffset_maxval_pos(self):
|
|
with self.assertRaisesRegex(Exception, 'Invalid Offset Specified'):
|
|
tup = self.s.vsGetFieldByOffset(0xffffffff)
|
|
|
|
def test_vsGetFieldByOffset_maxval_plus1_pos(self):
|
|
with self.assertRaisesRegex(Exception, 'Invalid Offset Specified'):
|
|
tup = self.s.vsGetFieldByOffset(len(self.s))
|
|
|
|
def test_vsGetFieldByOffset_nested1(self):
|
|
fname, fobj = self.s.vsGetFieldByOffset(5)
|
|
self.assertEqual(fname, 'two.foo')
|
|
self.assertIs(fobj, self.s.two.vsGetField('foo'))
|
|
|
|
def test_vsGetFieldByOffset_nested2(self):
|
|
fname, fobj = self.s.vsGetFieldByOffset(8)
|
|
self.assertEqual(fname, 'two.bar')
|
|
self.assertIs(fobj, self.s.two.vsGetField('bar'))
|
|
|
|
def test_vsGetFieldByOffset_nested3(self):
|
|
fname, fobj = self.s.vsGetFieldByOffset(12)
|
|
self.assertEqual(fname, 'two.baz')
|
|
self.assertIs(fobj, self.s.two.vsGetField('baz'))
|
|
|
|
def test_vsGetFieldByOffset_nested4(self):
|
|
fname, fobj = self.s.vsGetFieldByOffset(267)
|
|
self.assertEqual(fname, 'two.faz.alpha')
|
|
self.assertIs(fobj, self.s.two.faz.vsGetField('alpha'))
|
|
|
|
# TODO: could use envi.bits, but do we really want envi dep by default?
|
|
blkup = {}
|
|
bwidths = (8, 16, 24, 32, 64, )
|
|
for bwidth in bwidths:
|
|
umax = (2**bwidth) - 1
|
|
umin = 0
|
|
|
|
smax = (2**(bwidth-1)) - 1
|
|
smin = -(2**(bwidth-1))
|
|
|
|
blkup[bwidth] = (umin, umax, smin, smax)
|
|
|
|
class IntegerStruct(vstruct.VStruct):
|
|
def __init__(self):
|
|
vstruct.VStruct.__init__(self)
|
|
|
|
class VStructTypeTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.s = IntegerStruct()
|
|
|
|
def getTestFunc(name, vsval, val, expval):
|
|
|
|
def func(self):
|
|
self.s.vsAddField(name, vsval)
|
|
|
|
setattr(self.s, name, val)
|
|
|
|
rval = getattr(self.s, name)
|
|
self.assertEqual(rval, expval)
|
|
|
|
return func
|
|
|
|
tdefs = []
|
|
|
|
# dynamically generate the test definitions
|
|
# width, pname, vtype, test value, expected value
|
|
for bwidth in bwidths:
|
|
umin, umax, smin, smax = blkup[bwidth]
|
|
|
|
for ttype, mmin, mmax in ( ('u', umin, umax), ('', smin, smax), ):
|
|
pname = '{}int{}'.format(ttype, bwidth)
|
|
vtype = 'v_{}int{}'.format(ttype, bwidth)
|
|
vtype = getattr(p, vtype)
|
|
|
|
tup = (bwidth, pname, vtype, mmin, mmin)
|
|
tdefs.append(tup)
|
|
|
|
tup = (bwidth, pname, vtype, mmax, mmax)
|
|
tdefs.append(tup)
|
|
|
|
tup = (bwidth, pname, vtype, mmin-1, mmax)
|
|
tdefs.append(tup)
|
|
|
|
tup = (bwidth, pname, vtype, mmin-2, mmax-1)
|
|
tdefs.append(tup)
|
|
|
|
tup = (bwidth, pname, vtype, mmax+1, mmin)
|
|
tdefs.append(tup)
|
|
|
|
tup = (bwidth, pname, vtype, mmax+2, mmin+1)
|
|
tdefs.append(tup)
|
|
|
|
# generate unittest functions based on the test definitions
|
|
for width, pname, vtype, val, expval in tdefs:
|
|
tfunc = getTestFunc(pname, vtype(), val, expval)
|
|
tname = 'test_{}_{}_{}'.format(pname, val, expval)
|
|
setattr(VStructTypeTests, tname, tfunc)
|