remove dependency on warnings module
[p2pool.git] / p2pool / bitcoin / data.py
1 from __future__ import division
2
3 import hashlib
4 import struct
5
6 from . import base58, skiplists
7 from p2pool.util import bases, math, skiplist
8 import p2pool
9
10 class EarlyEnd(Exception):
11     pass
12
13 class LateEnd(Exception):
14     pass
15
16 def read((data, pos), length):
17     data2 = data[pos:pos + length]
18     if len(data2) != length:
19         raise EarlyEnd()
20     return data2, (data, pos + length)
21
22 def size((data, pos)):
23     return len(data) - pos
24
25 class Type(object):
26     # the same data can have only one unpacked representation, but multiple packed binary representations
27     
28     #def __hash__(self):
29     #    return hash(tuple(self.__dict__.items()))
30     
31     #def __eq__(self, other):
32     #    if not isinstance(other, Type):
33     #        raise NotImplementedError()
34     #    return self.__dict__ == other.__dict__
35     
36     def _unpack(self, data):
37         obj, (data2, pos) = self.read((data, 0))
38         
39         assert data2 is data
40         
41         if pos != len(data):
42             raise LateEnd()
43         
44         return obj
45     
46     def _pack(self, obj):
47         f = self.write(None, obj)
48         
49         res = []
50         while f is not None:
51             res.append(f[1])
52             f = f[0]
53         res.reverse()
54         return ''.join(res)
55     
56     
57     def unpack(self, data):
58         obj = self._unpack(data)
59         
60         if p2pool.DEBUG:
61             data2 = self._pack(obj)
62             if data2 != data:
63                 if self._unpack(data2) != obj:
64                     raise AssertionError()
65         
66         return obj
67     
68     def pack(self, obj):
69         data = self._pack(obj)
70         
71         if p2pool.DEBUG:
72             if self._unpack(data) != obj:
73                 raise AssertionError()
74         
75         return data
76     
77     
78     def pack_base58(self, obj):
79         return base58.base58_encode(self.pack(obj))
80     
81     def unpack_base58(self, base58_data):
82         return self.unpack(base58.base58_decode(base58_data))
83     
84     
85     def hash160(self, obj):
86         return ShortHashType().unpack(hashlib.new('ripemd160', hashlib.sha256(self.pack(obj)).digest()).digest())
87     
88     def hash256(self, obj):
89         return HashType().unpack(hashlib.sha256(hashlib.sha256(self.pack(obj)).digest()).digest())
90
91 class VarIntType(Type):
92     # redundancy doesn't matter here because bitcoin and p2pool both reencode before hashing
93     def read(self, file):
94         data, file = read(file, 1)
95         first = ord(data)
96         if first < 0xfd:
97             return first, file
98         elif first == 0xfd:
99             desc, length = '<H', 2
100         elif first == 0xfe:
101             desc, length = '<I', 4
102         elif first == 0xff:
103             desc, length = '<Q', 8
104         else:
105             raise AssertionError()
106         data, file = read(file, length)
107         return struct.unpack(desc, data)[0], file
108     
109     def write(self, file, item):
110         if item < 0xfd:
111             file = file, struct.pack('<B', item)
112         elif item <= 0xffff:
113             file = file, struct.pack('<BH', 0xfd, item)
114         elif item <= 0xffffffff:
115             file = file, struct.pack('<BI', 0xfe, item)
116         elif item <= 0xffffffffffffffff:
117             file = file, struct.pack('<BQ', 0xff, item)
118         else:
119             raise ValueError('int too large for varint')
120         return file
121
122 class VarStrType(Type):
123     _inner_size = VarIntType()
124     
125     def read(self, file):
126         length, file = self._inner_size.read(file)
127         return read(file, length)
128     
129     def write(self, file, item):
130         return self._inner_size.write(file, len(item)), item
131
132 class FixedStrType(Type):
133     def __init__(self, length):
134         self.length = length
135     
136     def read(self, file):
137         return read(file, self.length)
138     
139     def write(self, file, item):
140         if len(item) != self.length:
141             raise ValueError('incorrect length item!')
142         return file, item
143
144 class EnumType(Type):
145     def __init__(self, inner, values):
146         self.inner = inner
147         self.values = values
148         
149         self.keys = {}
150         for k, v in values.iteritems():
151             if v in self.keys:
152                 raise ValueError('duplicate value in values')
153             self.keys[v] = k
154     
155     def read(self, file):
156         data, file = self.inner.read(file)
157         return self.keys[data], file
158     
159     def write(self, file, item):
160         return self.inner.write(file, self.values[item])
161
162 class HashType(Type):
163     def read(self, file):
164         data, file = read(file, 256//8)
165         return int(data[::-1].encode('hex'), 16), file
166     
167     def write(self, file, item):
168         if not 0 <= item < 2**256:
169             raise ValueError('invalid hash value - %r' % (item,))
170         if item != 0 and item < 2**160:
171             print 'Very low hash value - maybe you meant to use ShortHashType? %x' % (item,)
172         return file, ('%064x' % (item,)).decode('hex')[::-1]
173
174 class ShortHashType(Type):
175     def read(self, file):
176         data, file = read(file, 160//8)
177         return int(data[::-1].encode('hex'), 16), file
178     
179     def write(self, file, item):
180         if not 0 <= item < 2**160:
181             raise ValueError('invalid hash value - %r' % (item,))
182         return file, ('%040x' % (item,)).decode('hex')[::-1]
183
184 class ListType(Type):
185     _inner_size = VarIntType()
186     
187     def __init__(self, type):
188         self.type = type
189     
190     def read(self, file):
191         length, file = self._inner_size.read(file)
192         res = []
193         for i in xrange(length):
194             item, file = self.type.read(file)
195             res.append(item)
196         return res, file
197     
198     def write(self, file, item):
199         file = self._inner_size.write(file, len(item))
200         for subitem in item:
201             file = self.type.write(file, subitem)
202         return file
203
204 class StructType(Type):
205     def __init__(self, desc):
206         self.desc = desc
207         self.length = struct.calcsize(self.desc)
208     
209     def read(self, file):
210         data, file = read(file, self.length)
211         res, = struct.unpack(self.desc, data)
212         return res, file
213     
214     def write(self, file, item):
215         data = struct.pack(self.desc, item)
216         if struct.unpack(self.desc, data)[0] != item:
217             # special test because struct doesn't error on some overflows
218             raise ValueError('''item didn't survive pack cycle (%r)''' % (item,))
219         return file, data
220
221 class IPV6AddressType(Type):
222     def read(self, file):
223         data, file = read(file, 16)
224         if data[:12] != '00000000000000000000ffff'.decode('hex'):
225             raise ValueError('ipv6 addresses not supported yet')
226         return '.'.join(str(ord(x)) for x in data[12:]), file
227     
228     def write(self, file, item):
229         bits = map(int, item.split('.'))
230         if len(bits) != 4:
231             raise ValueError('invalid address: %r' % (bits,))
232         data = '00000000000000000000ffff'.decode('hex') + ''.join(chr(x) for x in bits)
233         assert len(data) == 16, len(data)
234         return file, data
235
236 _record_types = {}
237
238 def get_record(fields):
239     fields = tuple(sorted(fields))
240     if 'keys' in fields:
241         raise ValueError()
242     if fields not in _record_types:
243         class _Record(object):
244             __slots__ = fields
245             def __getitem__(self, key):
246                 return getattr(self, key)
247             def __setitem__(self, key, value):
248                 setattr(self, key, value)
249             #def __iter__(self):
250             #    for field in self.__slots__:
251             #        yield field, getattr(self, field)
252             def keys(self):
253                 return self.__slots__
254             def __eq__(self, other):
255                 if isinstance(other, dict):
256                     return dict(self) == other
257                 elif isinstance(other, _Record):
258                     return all(self[k] == other[k] for k in self.keys())
259                 raise TypeError()
260             def __ne__(self, other):
261                 return not (self == other)
262         _record_types[fields] = _Record
263     return _record_types[fields]()
264
265 class ComposedType(Type):
266     def __init__(self, fields):
267         self.fields = fields
268     
269     def read(self, file):
270         item = get_record(k for k, v in self.fields)
271         for key, type_ in self.fields:
272             item[key], file = type_.read(file)
273         return item, file
274     
275     def write(self, file, item):
276         for key, type_ in self.fields:
277             file = type_.write(file, item[key])
278         return file
279
280 class ChecksummedType(Type):
281     def __init__(self, inner):
282         self.inner = inner
283     
284     def read(self, file):
285         obj, file = self.inner.read(file)
286         data = self.inner.pack(obj)
287         
288         checksum, file = read(file, 4)
289         if checksum != hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4]:
290             raise ValueError('invalid checksum')
291         
292         return obj, file
293     
294     def write(self, file, item):
295         data = self.inner.pack(item)
296         return (file, data), hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4]
297
298 class FloatingIntegerType(Type):
299     # redundancy doesn't matter here because bitcoin checks binary bits against its own computed bits
300     # so it will always be encoded 'normally' in blocks (they way bitcoin does it)
301     _inner = StructType('<I')
302     
303     def read(self, file):
304         bits, file = self._inner.read(file)
305         target = self._bits_to_target(bits)
306         if p2pool.DEBUG:
307             if self._target_to_bits(target) != bits:
308                 raise ValueError('bits in non-canonical form')
309         return target, file
310     
311     def write(self, file, item):
312         return self._inner.write(file, self._target_to_bits(item))
313     
314     def truncate_to(self, x):
315         return self._bits_to_target(self._target_to_bits(x, _check=False))
316     
317     def _bits_to_target(self, bits2):
318         target = math.shift_left(bits2 & 0x00ffffff, 8 * ((bits2 >> 24) - 3))
319         if p2pool.DEBUG:
320             assert target == self._bits_to_target1(struct.pack('<I', bits2))
321             assert self._target_to_bits(target, _check=False) == bits2
322         return target
323     
324     def _bits_to_target1(self, bits):
325         bits = bits[::-1]
326         length = ord(bits[0])
327         return bases.string_to_natural((bits[1:] + '\0'*length)[:length])
328     
329     def _target_to_bits(self, target, _check=True):
330         n = bases.natural_to_string(target)
331         if n and ord(n[0]) >= 128:
332             n = '\x00' + n
333         bits2 = (chr(len(n)) + (n + 3*chr(0))[:3])[::-1]
334         bits = struct.unpack('<I', bits2)[0]
335         if _check:
336             if self._bits_to_target(bits) != target:
337                 raise ValueError(repr((target, self._bits_to_target(bits, _check=False))))
338         return bits
339
340 class PossiblyNone(Type):
341     def __init__(self, none_value, inner):
342         self.none_value = none_value
343         self.inner = inner
344     
345     def read(self, file):
346         value, file = self.inner.read(file)
347         return None if value == self.none_value else value, file
348     
349     def write(self, file, item):
350         if item == self.none_value:
351             raise ValueError('none_value used')
352         return self.inner.write(file, self.none_value if item is None else item)
353
354 address_type = ComposedType([
355     ('services', StructType('<Q')),
356     ('address', IPV6AddressType()),
357     ('port', StructType('>H')),
358 ])
359
360 tx_type = ComposedType([
361     ('version', StructType('<I')),
362     ('tx_ins', ListType(ComposedType([
363         ('previous_output', PossiblyNone(dict(hash=0, index=2**32 - 1), ComposedType([
364             ('hash', HashType()),
365             ('index', StructType('<I')),
366         ]))),
367         ('script', VarStrType()),
368         ('sequence', PossiblyNone(2**32 - 1, StructType('<I'))),
369     ]))),
370     ('tx_outs', ListType(ComposedType([
371         ('value', StructType('<Q')),
372         ('script', VarStrType()),
373     ]))),
374     ('lock_time', StructType('<I')),
375 ])
376
377 block_header_type = ComposedType([
378     ('version', StructType('<I')),
379     ('previous_block', PossiblyNone(0, HashType())),
380     ('merkle_root', HashType()),
381     ('timestamp', StructType('<I')),
382     ('target', FloatingIntegerType()),
383     ('nonce', StructType('<I')),
384 ])
385
386 block_type = ComposedType([
387     ('header', block_header_type),
388     ('txs', ListType(tx_type)),
389 ])
390
391
392 merkle_record_type = ComposedType([
393     ('left', HashType()),
394     ('right', HashType()),
395 ])
396
397 def merkle_hash(tx_list):
398     if not tx_list:
399         return 0
400     hash_list = map(tx_type.hash256, tx_list)
401     while len(hash_list) > 1:
402         hash_list = [merkle_record_type.hash256(dict(left=left, right=left if right is None else right))
403             for left, right in zip(hash_list[::2], hash_list[1::2] + [None])]
404     return hash_list[0]
405
406 def target_to_average_attempts(target):
407     return 2**256//(target + 1)
408
409 # tx
410
411 def tx_get_sigop_count(tx):
412     return sum(script.get_sigop_count(txin['script']) for txin in tx['tx_ins']) + sum(script.get_sigop_count(txout['script']) for txout in tx['tx_outs'])
413
414 # human addresses
415
416 human_address_type = ChecksummedType(ComposedType([
417     ('version', StructType('<B')),
418     ('pubkey_hash', ShortHashType()),
419 ]))
420
421 pubkey_type = FixedStrType(65)
422
423 def pubkey_hash_to_address(pubkey_hash, net):
424     return human_address_type.pack_base58(dict(version=net.BITCOIN_ADDRESS_VERSION, pubkey_hash=pubkey_hash))
425
426 def pubkey_to_address(pubkey, net):
427     return pubkey_hash_to_address(pubkey_type.hash160(pubkey), net)
428
429 def address_to_pubkey_hash(address, net):
430     x = human_address_type.unpack_base58(address)
431     if x['version'] != net.BITCOIN_ADDRESS_VERSION:
432         raise ValueError('address not for this net!')
433     return x['pubkey_hash']
434
435 # transactions
436
437 def pubkey_to_script2(pubkey):
438     return ('\x41' + pubkey_type.pack(pubkey)) + '\xac'
439
440 def pubkey_hash_to_script2(pubkey_hash):
441     return '\x76\xa9' + ('\x14' + ShortHashType().pack(pubkey_hash)) + '\x88\xac'
442
443 # linked list tracker
444
445 class Tracker(object):
446     def __init__(self):
447         self.shares = {} # hash -> share
448         #self.ids = {} # hash -> (id, height)
449         self.reverse_shares = {} # previous_hash -> set of share_hashes
450         
451         self.heads = {} # head hash -> tail_hash
452         self.tails = {} # tail hash -> set of head hashes
453         
454         self.heights = {} # share_hash -> height_to, other_share_hash
455         
456         '''
457         self.id_generator = itertools.count()
458         self.tails_by_id = {}
459         '''
460         
461         self.get_nth_parent_hash = skiplists.DistanceSkipList(self)
462     
463     def add(self, share):
464         assert not isinstance(share, (int, long, type(None)))
465         if share.hash in self.shares:
466             return # XXX raise exception?
467         
468         '''
469         parent_id = self.ids.get(share.previous_hash, None)
470         children_ids = set(self.ids.get(share2_hash) for share2_hash in self.reverse_shares.get(share.hash, set()))
471         infos = set()
472         if parent_id is not None:
473             infos.add((parent_id[0], parent_id[1] + 1))
474         for child_id in children_ids:
475             infos.add((child_id[0], child_id[1] - 1))
476         if not infos:
477             infos.add((self.id_generator.next(), 0))
478         chosen = min(infos)
479         self.ids[share.hash] = chosen
480         '''
481         
482         self.shares[share.hash] = share
483         self.reverse_shares.setdefault(share.previous_hash, set()).add(share.hash)
484         
485         if share.hash in self.tails:
486             heads = self.tails.pop(share.hash)
487         else:
488             heads = set([share.hash])
489         
490         if share.previous_hash in self.heads:
491             tail = self.heads.pop(share.previous_hash)
492         else:
493             #dist, tail = self.get_height_and_last(share.previous_hash) # XXX this should be moved out of the critical area even though it shouldn't matter
494             tail = share.previous_hash
495             while tail in self.shares:
496                 tail = self.shares[tail].previous_hash
497         
498         self.tails.setdefault(tail, set()).update(heads)
499         if share.previous_hash in self.tails[tail]:
500             self.tails[tail].remove(share.previous_hash)
501         
502         for head in heads:
503             self.heads[head] = tail
504     
505     def test(self):
506         t = Tracker()
507         for s in self.shares.itervalues():
508             t.add(s)
509         
510         assert self.shares == t.shares, (self.shares, t.shares)
511         assert self.reverse_shares == t.reverse_shares, (self.reverse_shares, t.reverse_shares)
512         assert self.heads == t.heads, (self.heads, t.heads)
513         assert self.tails == t.tails, (self.tails, t.tails)
514     
515     def remove(self, share_hash):
516         assert isinstance(share_hash, (int, long, type(None)))
517         if share_hash not in self.shares:
518             raise KeyError()
519         share = self.shares[share_hash]
520         del share_hash
521         
522         if share.hash in self.heads and share.previous_hash in self.tails:
523             tail = self.heads.pop(share.hash)
524             self.tails[tail].remove(share.hash)
525             if not self.tails[share.previous_hash]:
526                 self.tails.pop(share.previous_hash)
527         elif share.hash in self.heads:
528             tail = self.heads.pop(share.hash)
529             self.tails[tail].remove(share.hash)
530             if self.reverse_shares[share.previous_hash] != set([share.hash]):
531                 pass # has sibling
532             else:
533                 self.tails[tail].add(share.previous_hash)
534                 self.heads[share.previous_hash] = tail
535         elif share.previous_hash in self.tails:
536             raise NotImplementedError() # will break other things..
537             heads = self.tails[share.previous_hash]
538             if len(self.reverse_shares[share.previous_hash]) > 1:
539                 raise NotImplementedError()
540             else:
541                 del self.tails[share.previous_hash]
542                 for head in heads:
543                     self.heads[head] = share.hash
544                 self.tails[share.hash] = set(heads)
545         else:
546             raise NotImplementedError()
547         
548         '''
549         height, tail = self.get_height_and_last(share.hash)
550         
551         if share.hash in self.heads:
552             my_heads = set([share.hash])
553         elif share.previous_hash in self.tails:
554             my_heads = self.tails[share.previous_hash]
555         else:
556             some_heads = self.tails[tail]
557             some_heads_heights = dict((that_head, self.get_height_and_last(that_head)[0]) for that_head in some_heads)
558             my_heads = set(that_head for that_head in some_heads
559                 if some_heads_heights[that_head] > height and
560                 self.get_nth_parent_hash(that_head, some_heads_heights[that_head] - height) == share.hash)
561         
562         if share.previous_hash != tail:
563             self.heads[share.previous_hash] = tail
564         
565         for head in my_heads:
566             if head != share.hash:
567                 self.heads[head] = share.hash
568             else:
569                 self.heads.pop(head)
570         
571         if share.hash in self.heads:
572             self.heads.pop(share.hash)
573         
574         
575         self.tails[tail].difference_update(my_heads)
576         if share.previous_hash != tail:
577             self.tails[tail].add(share.previous_hash)
578         if not self.tails[tail]:
579             self.tails.pop(tail)
580         if my_heads != set([share.hash]):
581             self.tails[share.hash] = set(my_heads) - set([share.hash])
582         '''
583         
584         self.shares.pop(share.hash)
585         self.reverse_shares[share.previous_hash].remove(share.hash)
586         if not self.reverse_shares[share.previous_hash]:
587             self.reverse_shares.pop(share.previous_hash)
588         
589         #assert self.test() is None
590     
591     def get_height(self, share_hash):
592         height, work, last = self.get_height_work_and_last(share_hash)
593         return height
594     
595     def get_work(self, share_hash):
596         height, work, last = self.get_height_work_and_last(share_hash)
597         return work
598     
599     def get_height_and_last(self, share_hash):
600         height, work, last = self.get_height_work_and_last(share_hash)
601         return height, last
602     
603     def get_height_work_and_last(self, share_hash):
604         assert isinstance(share_hash, (int, long, type(None)))
605         orig = share_hash
606         height = 0
607         work = 0
608         updates = []
609         while True:
610             if share_hash is None or share_hash not in self.shares:
611                 break
612             updates.append((share_hash, height, work))
613             if share_hash in self.heights:
614                 height_inc, share_hash, work_inc = self.heights[share_hash]
615             else:
616                 height_inc, share_hash, work_inc = 1, self.shares[share_hash].previous_hash, target_to_average_attempts(self.shares[share_hash].target)
617             height += height_inc
618             work += work_inc
619         for update_hash, height_then, work_then in updates:
620             self.heights[update_hash] = height - height_then, share_hash, work - work_then
621         return height, work, share_hash
622     
623     def get_chain_known(self, start_hash):
624         assert isinstance(start_hash, (int, long, type(None)))
625         '''
626         Chain starting with item of hash I{start_hash} of items that this Tracker contains
627         '''
628         item_hash_to_get = start_hash
629         while True:
630             if item_hash_to_get not in self.shares:
631                 break
632             share = self.shares[item_hash_to_get]
633             assert not isinstance(share, long)
634             yield share
635             item_hash_to_get = share.previous_hash
636     
637     def get_chain_to_root(self, start_hash, root=None):
638         assert isinstance(start_hash, (int, long, type(None)))
639         assert isinstance(root, (int, long, type(None)))
640         '''
641         Chain of hashes starting with share_hash of shares to the root (doesn't include root)
642         Raises an error if one is missing
643         '''
644         share_hash_to_get = start_hash
645         while share_hash_to_get != root:
646             share = self.shares[share_hash_to_get]
647             yield share
648             share_hash_to_get = share.previous_hash
649     
650     def get_best_hash(self):
651         '''
652         Returns hash of item with the most items in its chain
653         '''
654         if not self.heads:
655             return None
656         return max(self.heads, key=self.get_height_and_last)
657     
658     def get_highest_height(self):
659         return max(self.get_height_and_last(head)[0] for head in self.heads) if self.heads else 0
660
661 class FakeShare(object):
662     def __init__(self, **kwargs):
663         self.__dict__.update(kwargs)
664
665 if __name__ == '__main__':
666     
667     t = Tracker()
668     
669     for i in xrange(10000):
670         t.add(FakeShare(hash=i, previous_hash=i - 1 if i > 0 else None))
671     
672     #t.remove(99)
673     
674     print 'HEADS', t.heads
675     print 'TAILS', t.tails
676     
677     import random
678     
679     while False:
680         print
681         print '-'*30
682         print
683         t = Tracker()
684         for i in xrange(random.randrange(100)):
685             x = random.choice(list(t.shares) + [None])
686             print i, '->', x
687             t.add(FakeShare(i, x))
688         while t.shares:
689             x = random.choice(list(t.shares))
690             print 'DEL', x, t.__dict__
691             try:
692                 t.remove(x)
693             except NotImplementedError:
694                 print 'aborted; not implemented'
695         import time
696         time.sleep(.1)
697         print 'HEADS', t.heads
698         print 'TAILS', t.tails
699     
700     #for share_hash, share in sorted(t.shares.iteritems()):
701     #    print share_hash, share.previous_hash, t.heads.get(share_hash), t.tails.get(share_hash)
702     
703     #import sys;sys.exit()
704     
705     print t.get_nth_parent_hash(9000, 5000)
706     print t.get_nth_parent_hash(9001, 412)
707     #print t.get_nth_parent_hash(90, 51)
708     
709     for share_hash in sorted(t.shares):
710         print str(share_hash).rjust(4),
711         x = t.skips.get(share_hash, None)
712         if x is not None:
713             print str(x[0]).rjust(4),
714             for a in x[1]:
715                 print str(a).rjust(10),
716         print
717
718 # network definitions
719
720 class Mainnet(object):
721     BITCOIN_P2P_PREFIX = 'f9beb4d9'.decode('hex')
722     BITCOIN_P2P_PORT = 8333
723     BITCOIN_ADDRESS_VERSION = 0
724
725 class Testnet(object):
726     BITCOIN_P2P_PREFIX = 'fabfb5da'.decode('hex')
727     BITCOIN_P2P_PORT = 18333
728     BITCOIN_ADDRESS_VERSION = 111