From a733a1de2274b021237bbdd5f8053fb03a7dcac7 Mon Sep 17 00:00:00 2001 From: Forrest Voight Date: Sat, 3 Mar 2012 15:55:32 -0500 Subject: [PATCH] cleaned up ChecksummedType --- p2pool/bitcoin/data.py | 10 ++++++---- 1 files changed, 6 insertions(+), 4 deletions(-) diff --git a/p2pool/bitcoin/data.py b/p2pool/bitcoin/data.py index ad95ff3..481df58 100644 --- a/p2pool/bitcoin/data.py +++ b/p2pool/bitcoin/data.py @@ -12,22 +12,24 @@ def hash160(data): return pack.IntType(160).unpack(hashlib.new('ripemd160', hashlib.sha256(data).digest()).digest()) class ChecksummedType(pack.Type): - def __init__(self, inner): + def __init__(self, inner, checksum_func=lambda data: hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4]): self.inner = inner + self.checksum_func = checksum_func def read(self, file): obj, file = self.inner.read(file) data = self.inner.pack(obj) - checksum, file = pack.read(file, 4) - if checksum != hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4]: + calculated_checksum = self.checksum_func(data) + checksum, file = pack.read(file, len(calculated_checksum)) + if checksum != calculated_checksum: raise ValueError('invalid checksum') return obj, file def write(self, file, item): data = self.inner.pack(item) - return (file, data), hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4] + return (file, data), self.checksum_func(data) class FloatingInteger(object): __slots__ = ['bits', '_target'] -- 1.7.1