bitcointools/BCDataStream.py

   1 #
   2 # Workalike python implementation of Bitcoin's CDataStream class.
   3 #
   4 import struct
   5 import StringIO
   6 import mmap
   7
   8 class SerializationError(Exception):
   9   """ Thrown when there's a problem deserializing or serializing """
  10
  11 class BCDataStream(object):
  12   def __init__(self):
  13     self.input = None
  14     self.read_cursor = 0
  15
  16   def clear(self):
  17     self.input = None
  18     self.read_cursor = 0
  19
  20   def write(self, bytes):  # Initialize with string of bytes
  21     if self.input is None:
  22       self.input = bytes
  23     else:
  24       self.input += bytes
  25
  26   def map_file(self, file, start):  # Initialize with bytes from file
  27     self.input = mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ)
  28     self.read_cursor = start
  29   def seek_file(self, position):
  30     self.read_cursor = position
  31   def close_file(self):
  32     self.input.close()
  33
  34   def read_string(self):
  35     # Strings are encoded depending on length:
  36     # 0 to 252 :  1-byte-length followed by bytes (if any)
  37     # 253 to 65,535 : byte'253' 2-byte-length followed by bytes
  38     # 65,536 to 4,294,967,295 : byte '254' 4-byte-length followed by bytes
  39     # ... and the Bitcoin client is coded to understand:
  40     # greater than 4,294,967,295 : byte '255' 8-byte-length followed by bytes of string
  41     # ... but I don't think it actually handles any strings that big.
  42     if self.input is None:
  43       raise SerializationError("call write(bytes) before trying to deserialize")
  44
  45     try:
  46       length = self.read_compact_size()
  47     except IndexError:
  48       raise SerializationError("attempt to read past end of buffer")
  49
  50     return self.read_bytes(length)
  51
  52   def write_string(self, string):
  53     # Length-encoded as with read-string
  54     self.write_compact_size(len(string))
  55     self.write(string)
  56
  57   def read_bytes(self, length):
  58     try:
  59       result = self.input[self.read_cursor:self.read_cursor+length]
  60       self.read_cursor += length
  61       return result
  62     except IndexError:
  63       raise SerializationError("attempt to read past end of buffer")
  64
  65     return ''
  66
  67   def read_boolean(self): return self.read_bytes(1)[0] != chr(0)
  68   def read_int16(self): return self._read_num('<h')
  69   def read_uint16(self): return self._read_num('<H')
  70   def read_int32(self): return self._read_num('<i')
  71   def read_uint32(self): return self._read_num('<I')
  72   def read_int64(self): return self._read_num('<q')
  73   def read_uint64(self): return self._read_num('<Q')
  74
  75   def write_boolean(self, val): return self.write(chr(1) if val else chr(0))
  76   def write_int16(self, val): return self._write_num('<h', val)
  77   def write_uint16(self, val): return self._write_num('<H', val)
  78   def write_int32(self, val): return self._write_num('<i', val)
  79   def write_uint32(self, val): return self._write_num('<I', val)
  80   def write_int64(self, val): return self._write_num('<q', val)
  81   def write_uint64(self, val): return self._write_num('<Q', val)
  82
  83   def read_compact_size(self):
  84     size = ord(self.input[self.read_cursor])
  85     self.read_cursor += 1
  86     if size == 253:
  87       size = self._read_num('<H')
  88     elif size == 254:
  89       size = self._read_num('<I')
  90     elif size == 255:
  91       size = self._read_num('<Q')
  92     return size
  93
  94   def write_compact_size(self, size):
  95     if size < 0:
  96       raise SerializationError("attempt to write size < 0")
  97     elif size < 253:
  98        self.write(chr(size))
  99     elif size < 2**16:
 100       self.write('\xfd')
 101       self._write_num('<H', size)
 102     elif size < 2**32:
 103       self.write('\xfe')
 104       self._write_num('<I', size)
 105     elif size < 2**64:
 106       self.write('\xff')
 107       self._write_num('<Q', size)
 108
 109   def _read_num(self, format):
 110     (i,) = struct.unpack_from(format, self.input, self.read_cursor)
 111     self.read_cursor += struct.calcsize(format)
 112     return i
 113
 114   def _write_num(self, format, num):
 115     s = struct.pack(format, num)
 116     self.write(s)