added bitcoin.script.create_push_script and improved parsing

p2pool · Aug 8, 2012 · 0988632 · 0988632
1 parent 6bc81a7
commit 0988632
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 9 deletions.
diff --git a/p2pool/bitcoin/script.py b/p2pool/bitcoin/script.py
@@ -1,7 +1,7 @@
 from p2pool.util import math, pack
 
 def reads_nothing(f):
-    return '', f
+    return None, f
 def protoPUSH(length):
     return lambda f: pack.read(f, length)
 def protoPUSHDATA(size_len):
@@ -16,15 +16,15 @@ def _(f):
 for i in xrange(256):
     opcodes[i] = 'UNK_' + str(i), reads_nothing
 
-opcodes[0] = '0', reads_nothing
+opcodes[0] = 'PUSH', lambda f: ('', f)
 for i in xrange(1, 76):
-    opcodes[i] = 'PUSH%i' % i, protoPUSH(i)
-opcodes[76] = 'PUSHDATA1', protoPUSHDATA(1)
-opcodes[77] = 'PUSHDATA2', protoPUSHDATA(2)
-opcodes[78] = 'PUSHDATA4', protoPUSHDATA(4)
-opcodes[79] = '-1', reads_nothing
+    opcodes[i] = 'PUSH', protoPUSH(i)
+opcodes[76] = 'PUSH', protoPUSHDATA(1)
+opcodes[77] = 'PUSH', protoPUSHDATA(2)
+opcodes[78] = 'PUSH', protoPUSHDATA(4)
+opcodes[79] = 'PUSH', lambda f: ('\x81', f)
 for i in xrange(81, 97):
-    opcodes[i] = str(i - 80), reads_nothing
+    opcodes[i] = 'PUSH', lambda f, _i=i: (chr(_i - 80), f)
 
 opcodes[172] = 'CHECKSIG', reads_nothing
 opcodes[173] = 'CHECKSIGVERIFY', reads_nothing
@@ -48,3 +48,33 @@ def get_sigop_count(script):
         'CHECKMULTISIGVERIFY': 20,
     }
     return sum(weights.get(opcode_name, 0) for opcode_name, opcode_arg in parse(script))
+
+def create_push_script(datums): # datums can be ints or strs
+    res = []
+    for datum in datums:
+        if isinstance(datum, (int, long)):
+            if datum == -1 or 1 <= datum <= 16:
+                res.append(chr(datum + 80))
+                continue
+            negative = datum < 0
+            datum = math.natural_to_string(abs(datum))
+            if datum and ord(datum[0]) & 128:
+                datum = '\x00' + datum
+            if negative:
+                datum = chr(ord(datum[0]) + 128) + datum[1:]
+            datum = datum[::-1]
+        if len(datum) < 76:
+            res.append(chr(len(datum)))
+        elif len(datum) <= 0xff:
+            res.append(76)
+            res.append(chr(len(datum)))
+        elif len(datum) <= 0xffff:
+            res.append(77)
+            res.append(pack.IntType(16).pack(len(datum)))
+        elif len(datum) <= 0xffffffff:
+            res.append(78)
+            res.append(pack.IntType(32).pack(len(datum)))
+        else:
+            raise ValueError('string too long')
+        res.append(datum)
+    return ''.join(res)
diff --git a/p2pool/test/bitcoin/test_script.py b/p2pool/test/bitcoin/test_script.py
@@ -7,6 +7,6 @@ def test_all(self):
         data = '76  A9  14 89 AB CD EF AB BA AB BA AB BA AB BA AB BA AB BA AB BA AB BA  88 AC'.replace(' ', '').decode('hex')
         self.assertEquals(
             list(script.parse(data)),
-            [('UNK_118', ''), ('UNK_169', ''), ('PUSH20', '\x89\xab\xcd\xef\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba'), ('UNK_136', ''), ('CHECKSIG', '')],
+            [('UNK_118', None), ('UNK_169', None), ('PUSH', '\x89\xab\xcd\xef\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba\xab\xba'), ('UNK_136', None), ('CHECKSIG', None)],
         )
         self.assertEquals(script.get_sigop_count(data), 1)