tests/net_test/cstruct.py

   1 # Copyright 2014 The Android Open Source Project
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 """A simple module for declaring C-like structures.
  16
  17 Example usage:
  18
  19 >>> # Declare a struct type by specifying name, field formats and field names.
  20 ... # Field formats are the same as those used in the struct module.
  21 ... import cstruct
  22 >>> NLMsgHdr = cstruct.Struct("NLMsgHdr", "=LHHLL", "length type flags seq pid")
  23 >>>
  24 >>>
  25 >>> # Create instances from tuples or raw bytes. Data past the end is ignored.
  26 ... n1 = NLMsgHdr((44, 32, 0x2, 0, 491))
  27 >>> print n1
  28 NLMsgHdr(length=44, type=32, flags=2, seq=0, pid=491)
  29 >>>
  30 >>> n2 = NLMsgHdr("\x2c\x00\x00\x00\x21\x00\x02\x00"
  31 ...               "\x00\x00\x00\x00\xfe\x01\x00\x00" + "junk at end")
  32 >>> print n2
  33 NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510)
  34 >>>
  35 >>> # Serialize to raw bytes.
  36 ... print n1.Pack().encode("hex")
  37 2c0000002000020000000000eb010000
  38 >>>
  39 >>> # Parse the beginning of a byte stream as a struct, and return the struct
  40 ... # and the remainder of the stream for further reading.
  41 ... data = ("\x2c\x00\x00\x00\x21\x00\x02\x00"
  42 ...         "\x00\x00\x00\x00\xfe\x01\x00\x00"
  43 ...         "more data")
  44 >>> cstruct.Read(data, NLMsgHdr)
  45 (NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510), 'more data')
  46 >>>
  47 """
  48
  49 import ctypes
  50 import string
  51 import struct
  52
  53
  54 def CalcNumElements(fmt):
  55   size = struct.calcsize(fmt)
  56   elements = struct.unpack(fmt, "\x00" * size)
  57   return len(elements)
  58
  59
  60 def Struct(name, fmt, fieldnames, substructs={}):
  61   """Function that returns struct classes."""
  62
  63   class Meta(type):
  64
  65     def __len__(cls):
  66       return cls._length
  67
  68     def __init__(cls, unused_name, unused_bases, namespace):
  69       # Make the class object have the name that's passed in.
  70       type.__init__(cls, namespace["_name"], unused_bases, namespace)
  71
  72   class CStruct(object):
  73     """Class representing a C-like structure."""
  74
  75     __metaclass__ = Meta
  76
  77     # Name of the struct.
  78     _name = name
  79     # List of field names.
  80     _fieldnames = fieldnames
  81     # Dict mapping field indices to nested struct classes.
  82     _nested = {}
  83
  84     if isinstance(_fieldnames, str):
  85       _fieldnames = _fieldnames.split(" ")
  86
  87     # Parse fmt into _format, converting any S format characters to "XXs",
  88     # where XX is the length of the struct type's packed representation.
  89     _format = ""
  90     laststructindex = 0
  91     for i in xrange(len(fmt)):
  92       if fmt[i] == "S":
  93         # Nested struct. Record the index in our struct it should go into.
  94         index = CalcNumElements(fmt[:i])
  95         _nested[index] = substructs[laststructindex]
  96         laststructindex += 1
  97         _format += "%ds" % len(_nested[index])
  98       else:
  99          # Standard struct format character.
 100         _format += fmt[i]
 101
 102     _length = struct.calcsize(_format)
 103
 104     def _SetValues(self, values):
 105       super(CStruct, self).__setattr__("_values", list(values))
 106
 107     def _Parse(self, data):
 108       data = data[:self._length]
 109       values = list(struct.unpack(self._format, data))
 110       for index, value in enumerate(values):
 111         if isinstance(value, str) and index in self._nested:
 112           values[index] = self._nested[index](value)
 113       self._SetValues(values)
 114
 115     def __init__(self, values):
 116       # Initializing from a string.
 117       if isinstance(values, str):
 118         if len(values) < self._length:
 119           raise TypeError("%s requires string of length %d, got %d" %
 120                           (self._name, self._length, len(values)))
 121         self._Parse(values)
 122       else:
 123         # Initializing from a tuple.
 124         if len(values) != len(self._fieldnames):
 125           raise TypeError("%s has exactly %d fieldnames (%d given)" %
 126                           (self._name, len(self._fieldnames), len(values)))
 127         self._SetValues(values)
 128
 129     def _FieldIndex(self, attr):
 130       try:
 131         return self._fieldnames.index(attr)
 132       except ValueError:
 133         raise AttributeError("'%s' has no attribute '%s'" %
 134                              (self._name, attr))
 135
 136     def __getattr__(self, name):
 137       return self._values[self._FieldIndex(name)]
 138
 139     def __setattr__(self, name, value):
 140       self._values[self._FieldIndex(name)] = value
 141
 142     @classmethod
 143     def __len__(cls):
 144       return cls._length
 145
 146     def __ne__(self, other):
 147       return not self.__eq__(other)
 148
 149     def __eq__(self, other):
 150       return (isinstance(other, self.__class__) and
 151               self._name == other._name and
 152               self._fieldnames == other._fieldnames and
 153               self._values == other._values)
 154
 155     @staticmethod
 156     def _MaybePackStruct(value):
 157       if hasattr(value, "__metaclass__"):# and value.__metaclass__ == Meta:
 158         return value.Pack()
 159       else:
 160         return value
 161
 162     def Pack(self):
 163       values = [self._MaybePackStruct(v) for v in self._values]
 164       return struct.pack(self._format, *values)
 165
 166     def __str__(self):
 167       def FieldDesc(index, name, value):
 168         if isinstance(value, str) and any(
 169             c not in string.printable for c in value):
 170           value = value.encode("hex")
 171         return "%s=%s" % (name, value)
 172
 173       descriptions = [
 174           FieldDesc(i, n, v) for i, (n, v) in
 175           enumerate(zip(self._fieldnames, self._values))]
 176
 177       return "%s(%s)" % (self._name, ", ".join(descriptions))
 178
 179     def __repr__(self):
 180       return str(self)
 181
 182     def CPointer(self):
 183       """Returns a C pointer to the serialized structure."""
 184       buf = ctypes.create_string_buffer(self.Pack())
 185       # Store the C buffer in the object so it doesn't get garbage collected.
 186       super(CStruct, self).__setattr__("_buffer", buf)
 187       return ctypes.addressof(self._buffer)
 188
 189   return CStruct
 190
 191
 192 def Read(data, struct_type):
 193   length = len(struct_type)
 194   return struct_type(data), data[length:]