blob: c10667ad4bfda33f6b9709087245918996f160dd [file] [log] [blame]
# Copyright 2014 The Android Open Source Project
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
"""A simple module for declaring C-like structures.
Example usage:
>>> # Declare a struct type by specifying name, field formats and field names.
... # Field formats are the same as those used in the struct module, except:
... # - S: Nested Struct.
... # - A: NULL-padded ASCII string. Like s, but printing ignores contiguous
... # trailing NULL blocks at the end.
... import cstruct
>>> NLMsgHdr = cstruct.Struct("NLMsgHdr", "=LHHLL", "length type flags seq pid")
>>> # Create instances from a tuple of values, raw bytes, zero-initialized, or
>>> # using keywords.
... n1 = NLMsgHdr((44, 32, 0x2, 0, 491))
>>> print(n1)
NLMsgHdr(length=44, type=32, flags=2, seq=0, pid=491)
>>> n2 = NLMsgHdr("\x2c\x00\x00\x00\x21\x00\x02\x00"
... "\x00\x00\x00\x00\xfe\x01\x00\x00" + "junk at end")
>>> print(n2)
NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510)
>>> n3 = netlink.NLMsgHdr() # Zero-initialized
>>> print(n3)
NLMsgHdr(length=0, type=0, flags=0, seq=0, pid=0)
>>> n4 = netlink.NLMsgHdr(length=44, type=33) # Other fields zero-initialized
>>> print(n4)
NLMsgHdr(length=44, type=33, flags=0, seq=0, pid=0)
>>> # Serialize to raw bytes.
... print(n1.Pack().encode("hex"))
>>> # Parse the beginning of a byte stream as a struct, and return the struct
... # and the remainder of the stream for further reading.
... data = ("\x2c\x00\x00\x00\x21\x00\x02\x00"
... "\x00\x00\x00\x00\xfe\x01\x00\x00"
... "more data")
>>> cstruct.Read(data, NLMsgHdr)
(NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510), 'more data')
>>> # Structs can contain one or more nested structs. The nested struct types
... # are specified in a list as an optional last argument. Nested structs may
... # contain nested structs.
... S = cstruct.Struct("S", "=BI", "byte1 int2")
>>> N = cstruct.Struct("N", "!BSiS", "byte1 s2 int3 s2", [S, S])
>>> NN = cstruct.Struct("NN", "SHS", "s1 word2 n3", [S, N])
>>> nn = NN((S((1, 25000)), -29876, N((55, S((5, 6)), 1111, S((7, 8))))))
>>> nn.n3.s2.int2 = 5
import binascii
import ctypes
import string
import struct
import re
def _PythonFormat(fmt):
if "A" in fmt:
fmt = fmt.replace("A", "s")
return re.split('\d+$', fmt)[0]
def CalcSize(fmt):
return struct.calcsize(_PythonFormat(fmt))
def CalcNumElements(fmt):
fmt = _PythonFormat(fmt)
prevlen = len(fmt)
fmt = fmt.replace("S", "")
numstructs = prevlen - len(fmt)
size = struct.calcsize(fmt)
elements = struct.unpack(fmt, b"\x00" * size)
return len(elements) + numstructs
class StructMetaclass(type):
def __len__(cls):
return cls._length
def __init__(cls, unused_name, unused_bases, namespace):
# Make the class object have the name that's passed in.
type.__init__(cls, namespace["_name"], unused_bases, namespace)
def Struct(name, fmt, fieldnames, substructs={}):
"""Function that returns struct classes."""
# Hack to make struct classes use the StructMetaclass class on both python2 and
# python3. This is needed because in python2 the metaclass is assigned in the
# class definition, but in python3 it's passed into the constructor via
# keyword argument. Works by making all structs subclass CStructSuperclass,
# whose __new__ method uses StructMetaclass as its metaclass.
# A better option would be to use six.with_metaclass, but the existing python2
# VM image doesn't have the six module.
CStructSuperclass = type.__new__(StructMetaclass, 'unused', (), {})
class CStruct(CStructSuperclass):
"""Class representing a C-like structure."""
# Name of the struct.
_name = name
# List of field names.
_fieldnames = fieldnames
# Dict mapping field indices to nested struct classes.
_nested = {}
# List of string fields that are ASCII strings.
_asciiz = set()
_fieldnames = _fieldnames.split(" ")
# Parse fmt into _format, converting any S format characters to "XXs",
# where XX is the length of the struct type's packed representation.
_format = ""
laststructindex = 0
for i in range(len(fmt)):
if fmt[i] == "S":
# Nested struct. Record the index in our struct it should go into.
index = CalcNumElements(fmt[:i])
_nested[index] = substructs[laststructindex]
laststructindex += 1
_format += "%ds" % len(_nested[index])
elif fmt[i] == "A":
# Null-terminated ASCII string. Remove digits before the A, so we don't
# call CalcNumElements on an (invalid) format that ends with a digit.
start = i
while start > 0 and fmt[start - 1].isdigit(): start -= 1
index = CalcNumElements(fmt[:start])
_format += "s"
# Standard struct format character.
_format += fmt[i]
_length = CalcSize(_format)
offset_list = [0]
last_offset = 0
for i in range(len(_format)):
offset = CalcSize(_format[:i])
if offset > last_offset:
last_offset = offset
# A dictionary that maps field names to their offsets in the struct.
_offsets = dict(list(zip(_fieldnames, offset_list)))
# Check that the number of field names matches the number of fields.
numfields = len(struct.unpack(_format, b"\x00" * _length))
if len(_fieldnames) != numfields:
raise ValueError("Invalid cstruct: \"%s\" has %d elements, \"%s\" has %d."
% (fmt, numfields, fieldnames, len(_fieldnames)))
def _SetValues(self, values):
# Replace self._values with the given list. We can't do direct assignment
# because of the __setattr__ overload on this class.
super(CStruct, self).__setattr__("_values", list(values))
def _Parse(self, data):
data = data[:self._length]
values = list(struct.unpack(self._format, data))
for index, value in enumerate(values):
if isinstance(value, bytes) and index in self._nested:
values[index] = self._nested[index](value)
def __init__(self, tuple_or_bytes=None, **kwargs):
"""Construct an instance of this Struct.
1. With no args, the whole struct is zero-initialized.
2. With keyword args, the matching fields are populated; rest are zeroed.
3. With one tuple as the arg, the fields are assigned based on position.
4. With one bytes arg, the Struct is parsed from bytes.
if tuple_or_bytes and kwargs:
raise TypeError(
"%s: cannot specify both a tuple and keyword args" % self._name)
if tuple_or_bytes is None:
# Default construct from null bytes.
self._Parse(b"\x00" * len(self))
# If any keywords were supplied, set those fields.
for k, v in kwargs.items():
setattr(self, k, v)
elif isinstance(tuple_or_bytes, bytes):
# Initializing from bytes.
if len(tuple_or_bytes) < self._length:
raise TypeError("%s requires a bytes object of length %d, got %d" %
(self._name, self._length, len(tuple_or_bytes)))
# Initializing from a tuple.
if len(tuple_or_bytes) != len(self._fieldnames):
raise TypeError("%s has exactly %d fieldnames: (%s), %d given: (%s)" %
(self._name, len(self._fieldnames),
", ".join(self._fieldnames), len(tuple_or_bytes),
", ".join(str(x) for x in tuple_or_bytes)))
def _FieldIndex(self, attr):
return self._fieldnames.index(attr)
except ValueError:
raise AttributeError("'%s' has no attribute '%s'" %
(self._name, attr))
def __getattr__(self, name):
return self._values[self._FieldIndex(name)]
def __setattr__(self, name, value):
# TODO: check value type against self._format and throw here, or else
# callers get an unhelpful exception when they call Pack().
self._values[self._FieldIndex(name)] = value
def offset(self, name):
if "." in name:
raise NotImplementedError("offset() on nested field")
return self._offsets[name]
def __len__(cls):
return cls._length
def __ne__(self, other):
return not self.__eq__(other)
def __eq__(self, other):
return (isinstance(other, self.__class__) and
self._name == other._name and
self._fieldnames == other._fieldnames and
self._values == other._values)
def _MaybePackStruct(value):
if isinstance(type(value), StructMetaclass):
return value.Pack()
return value
def Pack(self):
values = [self._MaybePackStruct(v) for v in self._values]
return struct.pack(self._format, *values)
def __str__(self):
def HasNonPrintableChar(s):
for c in s:
# Iterating over bytes yields chars in python2 but ints in python3.
if isinstance(c, int): c = chr(c)
if c not in string.printable: return True
return False
def FieldDesc(index, name, value):
if isinstance(value, bytes) or isinstance(value, str):
if index in self._asciiz:
# TODO: use "backslashreplace" when python 2 is no longer supported.
value = value.rstrip(b"\x00").decode(errors="ignore")
elif HasNonPrintableChar(value):
value = binascii.hexlify(value).decode()
return "%s=%s" % (name, str(value))
descriptions = [
FieldDesc(i, n, v) for i, (n, v) in
enumerate(zip(self._fieldnames, self._values))]
return "%s(%s)" % (self._name, ", ".join(descriptions))
def __repr__(self):
return str(self)
def CPointer(self):
"""Returns a C pointer to the serialized structure."""
buf = ctypes.create_string_buffer(self.Pack())
# Store the C buffer in the object so it doesn't get garbage collected.
super(CStruct, self).__setattr__("_buffer", buf)
return ctypes.addressof(self._buffer)
return CStruct
def Read(data, struct_type):
length = len(struct_type)
return struct_type(data), data[length:]