p2pool/util/pack.py

   1 import binascii
   2 import struct
   3
   4 import p2pool
   5
   6 class EarlyEnd(Exception):
   7     pass
   8
   9 class LateEnd(Exception):
  10     pass
  11
  12 def read((data, pos), length):
  13     data2 = data[pos:pos + length]
  14     if len(data2) != length:
  15         raise EarlyEnd()
  16     return data2, (data, pos + length)
  17
  18 def size((data, pos)):
  19     return len(data) - pos
  20
  21 class Type(object):
  22     __slots__ = []
  23
  24     def __hash__(self):
  25         rval = getattr(self, '_hash', None)
  26         if rval is None:
  27             try:
  28                 rval = self._hash = hash((type(self), frozenset(self.__dict__.items())))
  29             except:
  30                 print self.__dict__
  31                 raise
  32         return rval
  33
  34     def __eq__(self, other):
  35         return type(other) is type(self) and other.__dict__ == self.__dict__
  36
  37     def __ne__(self, other):
  38         return not (self == other)
  39
  40     def _unpack(self, data):
  41         obj, (data2, pos) = self.read((data, 0))
  42
  43         assert data2 is data
  44
  45         if pos != len(data):
  46             raise LateEnd()
  47
  48         return obj
  49
  50     def _pack(self, obj):
  51         f = self.write(None, obj)
  52
  53         res = []
  54         while f is not None:
  55             res.append(f[1])
  56             f = f[0]
  57         res.reverse()
  58         return ''.join(res)
  59
  60
  61     def unpack(self, data):
  62         obj = self._unpack(data)
  63
  64         if p2pool.DEBUG:
  65             if self._pack(obj) != data:
  66                 raise AssertionError()
  67
  68         return obj
  69
  70     def pack(self, obj):
  71         data = self._pack(obj)
  72
  73         if p2pool.DEBUG:
  74             if self._unpack(data) != obj:
  75                 raise AssertionError((self._unpack(data), obj))
  76
  77         return data
  78
  79     def packed_size(self, obj):
  80         if hasattr(obj, '_packed_size') and obj._packed_size is not None:
  81             type_obj, packed_size = obj._packed_size
  82             if type_obj is self:
  83                 return packed_size
  84
  85         packed_size = len(self.pack(obj))
  86
  87         if hasattr(obj, '_packed_size'):
  88             obj._packed_size = self, packed_size
  89
  90         return packed_size
  91
  92 class VarIntType(Type):
  93     def read(self, file):
  94         data, file = read(file, 1)
  95         first = ord(data)
  96         if first < 0xfd:
  97             return first, file
  98         if first == 0xfd:
  99             desc, length, minimum = '<H', 2, 0xfd
 100         elif first == 0xfe:
 101             desc, length, minimum = '<I', 4, 2**16
 102         elif first == 0xff:
 103             desc, length, minimum = '<Q', 8, 2**32
 104         else:
 105             raise AssertionError()
 106         data2, file = read(file, length)
 107         res, = struct.unpack(desc, data2)
 108         if res < minimum:
 109             raise AssertionError('VarInt not canonically packed')
 110         return res, file
 111
 112     def write(self, file, item):
 113         if item < 0xfd:
 114             return file, struct.pack('<B', item)
 115         elif item <= 0xffff:
 116             return file, struct.pack('<BH', 0xfd, item)
 117         elif item <= 0xffffffff:
 118             return file, struct.pack('<BI', 0xfe, item)
 119         elif item <= 0xffffffffffffffff:
 120             return file, struct.pack('<BQ', 0xff, item)
 121         else:
 122             raise ValueError('int too large for varint')
 123
 124 class VarStrType(Type):
 125     _inner_size = VarIntType()
 126
 127     def read(self, file):
 128         length, file = self._inner_size.read(file)
 129         return read(file, length)
 130
 131     def write(self, file, item):
 132         return self._inner_size.write(file, len(item)), item
 133
 134 class EnumType(Type):
 135     def __init__(self, inner, pack_to_unpack):
 136         self.inner = inner
 137         self.pack_to_unpack = pack_to_unpack
 138
 139         self.unpack_to_pack = {}
 140         for k, v in pack_to_unpack.iteritems():
 141             if v in self.unpack_to_pack:
 142                 raise ValueError('duplicate value in pack_to_unpack')
 143             self.unpack_to_pack[v] = k
 144
 145     def read(self, file):
 146         data, file = self.inner.read(file)
 147         if data not in self.pack_to_unpack:
 148             raise ValueError('enum data (%r) not in pack_to_unpack (%r)' % (data, self.pack_to_unpack))
 149         return self.pack_to_unpack[data], file
 150
 151     def write(self, file, item):
 152         if item not in self.unpack_to_pack:
 153             raise ValueError('enum item (%r) not in unpack_to_pack (%r)' % (item, self.unpack_to_pack))
 154         return self.inner.write(file, self.unpack_to_pack[item])
 155
 156 class ListType(Type):
 157     _inner_size = VarIntType()
 158
 159     def __init__(self, type):
 160         self.type = type
 161
 162     def read(self, file):
 163         length, file = self._inner_size.read(file)
 164         res = [None]*length
 165         for i in xrange(length):
 166             res[i], file = self.type.read(file)
 167         return res, file
 168
 169     def write(self, file, item):
 170         file = self._inner_size.write(file, len(item))
 171         for subitem in item:
 172             file = self.type.write(file, subitem)
 173         return file
 174
 175 class StructType(Type):
 176     __slots__ = 'desc length'.split(' ')
 177
 178     def __init__(self, desc):
 179         self.desc = desc
 180         self.length = struct.calcsize(self.desc)
 181
 182     def read(self, file):
 183         data, file = read(file, self.length)
 184         return struct.unpack(self.desc, data)[0], file
 185
 186     def write(self, file, item):
 187         return file, struct.pack(self.desc, item)
 188
 189 class IntType(Type):
 190     __slots__ = 'bytes step format_str max'.split(' ')
 191
 192     def __new__(cls, bits, endianness='little'):
 193         assert bits % 8 == 0
 194         assert endianness in ['little', 'big']
 195         if bits in [8, 16, 32, 64]:
 196             return StructType(('<' if endianness == 'little' else '>') + {8: 'B', 16: 'H', 32: 'I', 64: 'Q'}[bits])
 197         else:
 198             return Type.__new__(cls, bits, endianness)
 199
 200     def __init__(self, bits, endianness='little'):
 201         assert bits % 8 == 0
 202         assert endianness in ['little', 'big']
 203         self.bytes = bits//8
 204         self.step = -1 if endianness == 'little' else 1
 205         self.format_str = '%%0%ix' % (2*self.bytes)
 206         self.max = 2**bits
 207
 208     def read(self, file, b2a_hex=binascii.b2a_hex):
 209         if self.bytes == 0:
 210             return 0, file
 211         data, file = read(file, self.bytes)
 212         return int(b2a_hex(data[::self.step]), 16), file
 213
 214     def write(self, file, item, a2b_hex=binascii.a2b_hex):
 215         if self.bytes == 0:
 216             return file
 217         if not 0 <= item < self.max:
 218             raise ValueError('invalid int value - %r' % (item,))
 219         return file, a2b_hex(self.format_str % (item,))[::self.step]
 220
 221 class IPV6AddressType(Type):
 222     def read(self, file):
 223         data, file = read(file, 16)
 224         if data[:12] == '00000000000000000000ffff'.decode('hex'):
 225             return '.'.join(str(ord(x)) for x in data[12:]), file
 226         return ':'.join(data[i*2:(i+1)*2].encode('hex') for i in xrange(8)), file
 227
 228     def write(self, file, item):
 229         if ':' in item:
 230             data = ''.join(item.replace(':', '')).decode('hex')
 231         else:
 232             bits = map(int, item.split('.'))
 233             if len(bits) != 4:
 234                 raise ValueError('invalid address: %r' % (bits,))
 235             data = '00000000000000000000ffff'.decode('hex') + ''.join(chr(x) for x in bits)
 236         assert len(data) == 16, len(data)
 237         return file, data
 238
 239 _record_types = {}
 240
 241 def get_record(fields):
 242     fields = tuple(sorted(fields))
 243     if 'keys' in fields or '_packed_size' in fields:
 244         raise ValueError()
 245     if fields not in _record_types:
 246         class _Record(object):
 247             __slots__ = fields + ('_packed_size',)
 248             def __init__(self):
 249                 self._packed_size = None
 250             def __repr__(self):
 251                 return repr(dict(self))
 252             def __getitem__(self, key):
 253                 return getattr(self, key)
 254             def __setitem__(self, key, value):
 255                 setattr(self, key, value)
 256             #def __iter__(self):
 257             #    for field in fields:
 258             #        yield field, getattr(self, field)
 259             def keys(self):
 260                 return fields
 261             def get(self, key, default=None):
 262                 return getattr(self, key, default)
 263             def __eq__(self, other):
 264                 if isinstance(other, dict):
 265                     return dict(self) == other
 266                 elif isinstance(other, _Record):
 267                     for k in fields:
 268                         if getattr(self, k) != getattr(other, k):
 269                             return False
 270                     return True
 271                 elif other is None:
 272                     return False
 273                 raise TypeError()
 274             def __ne__(self, other):
 275                 return not (self == other)
 276         _record_types[fields] = _Record
 277     return _record_types[fields]
 278
 279 class ComposedType(Type):
 280     def __init__(self, fields):
 281         self.fields = list(fields)
 282         self.field_names = set(k for k, v in fields)
 283         self.record_type = get_record(k for k, v in self.fields)
 284
 285     def read(self, file):
 286         item = self.record_type()
 287         for key, type_ in self.fields:
 288             item[key], file = type_.read(file)
 289         return item, file
 290
 291     def write(self, file, item):
 292         assert set(item.keys()) == self.field_names, (set(item.keys()) - self.field_names, self.field_names - set(item.keys()))
 293         for key, type_ in self.fields:
 294             file = type_.write(file, item[key])
 295         return file
 296
 297 class PossiblyNoneType(Type):
 298     def __init__(self, none_value, inner):
 299         self.none_value = none_value
 300         self.inner = inner
 301
 302     def read(self, file):
 303         value, file = self.inner.read(file)
 304         return None if value == self.none_value else value, file
 305
 306     def write(self, file, item):
 307         if item == self.none_value:
 308             raise ValueError('none_value used')
 309         return self.inner.write(file, self.none_value if item is None else item)
 310
 311 class FixedStrType(Type):
 312     def __init__(self, length):
 313         self.length = length
 314
 315     def read(self, file):
 316         return read(file, self.length)
 317
 318     def write(self, file, item):
 319         if len(item) != self.length:
 320             raise ValueError('incorrect length item!')
 321         return file, item