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