From b85543cb7361d6ba27c0eeca756eec9fd5395b36 Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter@wuille.net>
Date: Tue, 26 May 2020 14:03:00 -0700
Subject: [PATCH] test: Add Python MuHash3072 implementation to test framework

---
 test/functional/test_framework/muhash.py | 90 ++++++++++++++++++++++++
 1 file changed, 90 insertions(+)
 create mode 100644 test/functional/test_framework/muhash.py

diff --git a/test/functional/test_framework/muhash.py b/test/functional/test_framework/muhash.py
new file mode 100644
index 00000000000..7dfce1b7565
--- /dev/null
+++ b/test/functional/test_framework/muhash.py
@@ -0,0 +1,90 @@
+# Copyright (c) 2020 Pieter Wuille
+# Distributed under the MIT software license, see the accompanying
+# file COPYING or http://www.opensource.org/licenses/mit-license.php.
+"""Native Python MuHash3072 implementation."""
+
+import hashlib
+
+from .util import modinv
+
+def rot32(v, bits):
+    """Rotate the 32-bit value v left by bits bits."""
+    bits %= 32  # Make sure the term below does not throw an exception
+    return ((v << bits) & 0xffffffff) | (v >> (32 - bits))
+
+def chacha20_doubleround(s):
+    """Apply a ChaCha20 double round to 16-element state array s.
+
+    See https://cr.yp.to/chacha/chacha-20080128.pdf and https://tools.ietf.org/html/rfc8439
+    """
+    QUARTER_ROUNDS = [(0, 4, 8, 12),
+                      (1, 5, 9, 13),
+                      (2, 6, 10, 14),
+                      (3, 7, 11, 15),
+                      (0, 5, 10, 15),
+                      (1, 6, 11, 12),
+                      (2, 7, 8, 13),
+                      (3, 4, 9, 14)]
+
+    for a, b, c, d in QUARTER_ROUNDS:
+        s[a] = (s[a] + s[b]) & 0xffffffff
+        s[d] = rot32(s[d] ^ s[a], 16)
+        s[c] = (s[c] + s[d]) & 0xffffffff
+        s[b] = rot32(s[b] ^ s[c], 12)
+        s[a] = (s[a] + s[b]) & 0xffffffff
+        s[d] = rot32(s[d] ^ s[a], 8)
+        s[c] = (s[c] + s[d]) & 0xffffffff
+        s[b] = rot32(s[b] ^ s[c], 7)
+
+def chacha20_32_to_384(key32):
+    """Specialized ChaCha20 implementation with 32-byte key, 0 IV, 384-byte output."""
+    # See RFC 8439 section 2.3 for chacha20 parameters
+    CONSTANTS = [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574]
+
+    key_bytes = [0]*8
+    for i in range(8):
+        key_bytes[i] = int.from_bytes(key32[(4 * i):(4 * (i+1))], 'little')
+
+    INITIALIZATION_VECTOR = [0] * 4
+    init = CONSTANTS + key_bytes + INITIALIZATION_VECTOR
+    out = bytearray()
+    for counter in range(6):
+        init[12] = counter
+        s = init.copy()
+        for _ in range(10):
+            chacha20_doubleround(s)
+        for i in range(16):
+            out.extend(((s[i] + init[i]) & 0xffffffff).to_bytes(4, 'little'))
+    return bytes(out)
+
+def data_to_num3072(data):
+    """Hash a 32-byte array data to a 3072-bit number using 6 Chacha20 operations."""
+    bytes384 = chacha20_32_to_384(data)
+    return int.from_bytes(bytes384, 'little')
+
+class MuHash3072:
+    """Class representing the MuHash3072 computation of a set.
+
+    See https://cseweb.ucsd.edu/~mihir/papers/inchash.pdf and https://lists.linuxfoundation.org/pipermail/bitcoin-dev/2017-May/014337.html
+    """
+
+    MODULUS = 2**3072 - 1103717
+
+    def __init__(self):
+        """Initialize for an empty set."""
+        self.numerator = 1
+        self.denominator = 1
+
+    def insert(self, data):
+        """Insert a byte array data in the set."""
+        self.numerator = (self.numerator * data_to_num3072(data)) % self.MODULUS
+
+    def remove(self, data):
+        """Remove a byte array from the set."""
+        self.denominator = (self.denominator * data_to_num3072(data)) % self.MODULUS
+
+    def digest(self):
+        """Extract the final hash. Does not modify this object."""
+        val = (self.numerator * modinv(self.denominator, self.MODULUS)) % self.MODULUS
+        bytes384 = val.to_bytes(384, 'little')
+        return hashlib.sha256(bytes384).digest()