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 17Example 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 28NLMsgHdr(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 33NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510) 34>>> 35>>> # Serialize to raw bytes. 36... print n1.Pack().encode("hex") 372c0000002000020000000000eb010000 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 49import ctypes 50import string 51import struct 52 53 54def CalcNumElements(fmt): 55 size = struct.calcsize(fmt) 56 elements = struct.unpack(fmt, "\x00" * size) 57 return len(elements) 58 59 60def 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 192def Read(data, struct_type): 193 length = len(struct_type) 194 return struct_type(data), data[length:] 195