Skip to content

Instantly share code, notes, and snippets.

@dev-zzo
Last active August 22, 2022 15:40
Show Gist options
  • Save dev-zzo/31eef14df764040b1081882284b79305 to your computer and use it in GitHub Desktop.
Save dev-zzo/31eef14df764040b1081882284b79305 to your computer and use it in GitHub Desktop.
Simple but better structs for Python :)
import struct
class BetterStructMeta(type):
def __new__(cls, clsname, superclasses, attributedict):
if clsname != 'BetterStruct':
fields = attributedict['__fields__']
field_types = [ _[0] for _ in fields ]
field_names = [ _[1] for _ in fields if _[1] is not None ]
size = 0
for fielddef in fields:
t = fielddef[0]
n = fielddef[1]
if isinstance(t, str):
size += struct.Struct(t).size
else:
size += t.size
attributedict['__names__'] = field_names
attributedict['size'] = size
return type.__new__(cls, clsname, superclasses, attributedict)
class BetterStruct(metaclass=BetterStructMeta):
def __init__(self, **kwds):
for fielddef in self.__fields__:
t = fielddef[0]
n = fielddef[1]
if n is None:
continue
setattr(self, n, None)
if isinstance(t, str):
if t[-1] == 's':
setattr(self, n, b'')
elif t in ('Q', 'I', 'H', 'B'):
setattr(self, n, 0)
for n, v in kwds.items():
if n not in self.__names__:
raise KeyError("attribute name '%s' unknown" % n)
setattr(self, n, v)
@classmethod
def unpack_from(cls, buffer, offset=0):
if len(buffer) - offset < cls.size:
raise ValueError("unpack_from requires a buffer of at least %d bytes; %d available" % (cls.size, len(buffer) - offset))
fields = []
for fielddef in cls.__fields__:
t = fielddef[0]
n = fielddef[1]
if isinstance(t, str):
s = struct.Struct(t)
size = s.size
if n is not None:
value = s.unpack_from(buffer, offset)[0]
if t[-1] == 's' and len(fielddef) > 2 and fielddef[2]:
value = value.rstrip(b"\0")
fields.append(value)
else:
size = t.size
if n is not None:
fields.append(t.unpack_from(buffer, offset))
offset += size
instance = cls()
for n, v in zip(cls.__names__, fields):
setattr(instance, n, v)
try:
instance.validate()
except AttributeError:
pass
return instance
@classmethod
def read_from(cls, fp):
data = fp.read(cls.size)
if not data:
return None
if len(data) < cls.size:
raise ValueError("insufficient data read from the file object")
return cls.unpack_from(data)
def pack(self):
data = bytearray()
for fielddef in self.__fields__:
t = fielddef[0]
n = fielddef[1]
if n is not None:
v = getattr(self, n)
if isinstance(t, str):
data.extend(struct.Struct(t).pack(v))
else:
data.extend(v.pack())
else:
data.extend(struct.Struct(t).pack())
return data
def write(self, fp):
fp.write(self.pack())
def __str__(self):
return '(' + ', '.join([ "'%s': %s" % (n, repr(getattr(self, n))) for n in self.__names__ if n is not None ]) + ')'
def __repr__(self):
return str(self)
# Example structure definition
class ImageHeader(BetterStruct):
__endian__ = '<'
__fields__ = [
('20s', 'hash'),
('4s', 'signature'),
('B', 'major_rev'),
('B', 'minor_rev'),
('H', 'flags'),
('I', 'total_blocks'),
('I', 'first_boot_tag_block'),
('I', 'first_boot_section_id'),
('H', 'key_count'),
('H', 'key_dict_block'),
('H', 'header_blocks'),
('H', 'section_count'),
('H', 'section_header_size'),
('2x', None), # for padding etc
]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment