Merge pull request #11 from str4d/sapling-note-encryption

Sapling note encryption test vectors

README.md

@@ -1,3 +1,5 @@
 Code to generate test vectors for various parts of Zcash.
 
 Requires `pyblake2`.
+
+`sapling_note_encryption.py` also requires `chacha20poly1305`.

sapling_note_encryption.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+import sys; assert sys.version_info[0] >= 3, "Python 3 required."
+
+from chacha20poly1305 import ChaCha20Poly1305
+import os
+from pyblake2 import blake2b
+import struct
+
+from sapling_generators import VALUE_COMMITMENT_VALUE_BASE, VALUE_COMMITMENT_RANDOMNESS_BASE
+from sapling_jubjub import Fr, JUBJUB_COFACTOR
+from sapling_key_components import SpendingKey, diversify_hash
+from sapling_notes import note_commit
+from sapling_utils import leos2bsp, leos2ip
+from tv_output import render_args, render_tv
+
+
+def kdf_sapling(shared_secret, epk):
+    digest = blake2b(digest_size=32, person=b'Zcash_SaplingKDF')
+    digest.update(bytes(shared_secret))
+    digest.update(bytes(epk))
+    return digest.digest()
+
+def prf_ock(ovk, cv, cmu, ephemeral_key):
+    digest = blake2b(digest_size=32, person=b'Zcash_Derive_ock')
+    digest.update(ovk)
+    digest.update(cv)
+    digest.update(cmu)
+    digest.update(ephemeral_key)
+    return digest.digest()
+
+class SaplingKeyAgreement(object):
+    @staticmethod
+    def private(random):
+        return Fr(leos2ip(random(32)))
+
+    @staticmethod
+    def derive_public(esk, g_d):
+        return g_d * esk
+
+    @staticmethod
+    def agree(esk, pk_d):
+        return pk_d * esk * JUBJUB_COFACTOR
+
+class SaplingSym(object):
+    @staticmethod
+    def k(random):
+        return random(32)
+
+    @staticmethod
+    def encrypt(key, plaintext):
+        cip = ChaCha20Poly1305(key)
+        return bytes(cip.encrypt(b'\x00' * 12, plaintext))
+
+
+class SaplingNotePlaintext(object):
+    def __init__(self, d, v, rcm, memo):
+        self.d = d
+        self.v = v
+        self.rcm = rcm
+        self.memo = memo
+
+    def __bytes__(self):
+        return (
+            b'\x01' +
+            self.d +
+            struct.pack('<Q', self.v) +
+            bytes(self.rcm) +
+            self.memo
+        )
+
+class SaplingNoteEncryption(object):
+    def __init__(self, random=os.urandom):
+        self._random = random
+
+    def encrypt(self, np, pk_d_new, g_d_new, cv_new, cm_new, ovk=None):
+        esk = SaplingKeyAgreement.private(self._random)
+        epk = SaplingKeyAgreement.derive_public(esk, g_d_new)
+        p_enc = bytes(np)
+        shared_secret = SaplingKeyAgreement.agree(esk, pk_d_new)
+        k_enc = kdf_sapling(shared_secret, epk)
+        c_enc = SaplingSym.encrypt(k_enc, p_enc)
+
+        if not ovk:
+            ock = SaplingSym.k(self._random)
+            op = self._random(64)
+        else:
+            cv = bytes(cv_new)
+            cmu = bytes(cm_new.u)
+            ephemeral_key = bytes(epk)
+            ock = prf_ock(ovk, cv, cmu, ephemeral_key)
+            op = bytes(pk_d_new) + bytes(esk)
+
+        c_out = SaplingSym.encrypt(ock, op)
+
+        return (esk, epk, shared_secret, k_enc, p_enc, c_enc, ock, op, c_out)
+
+
+def main():
+    args = render_args()
+
+    from random import Random
+    rng = Random(0xabad533d)
+    def randbytes(l):
+        ret = []
+        while len(ret) < l:
+            ret.append(rng.randrange(0, 256))
+        return bytes(ret)
+    ne = SaplingNoteEncryption(randbytes)
+
+    test_vectors = []
+    for i in range(0, 10):
+        sk = SpendingKey(bytes([i] * 32))
+        pk_d = sk.default_pkd()
+        g_d = diversify_hash(sk.default_d())
+
+        np = SaplingNotePlaintext(
+            sk.default_d(),
+            100000000 * (i+1),
+            Fr(8890123457840276890326754358439057438290574382905).exp(i+1),
+            b'\xf6' + b'\x00'*511)
+        cv = VALUE_COMMITMENT_VALUE_BASE * Fr(np.v) + VALUE_COMMITMENT_RANDOMNESS_BASE * np.rcm
+        cm = note_commit(
+            np.rcm,
+            leos2bsp(bytes(g_d)),
+            leos2bsp(bytes(pk_d)),
+            np.v)
+
+        (
+            esk, epk,
+            shared_secret,
+            k_enc, p_enc, c_enc,
+            ock, op, c_out,
+        ) = ne.encrypt(np, pk_d, g_d, cv, cm, sk.ovk())
+
+        test_vectors.append({
+            'ovk': sk.ovk(),
+            'ivk': bytes(sk.ivk()),
+            'default_d': sk.default_d(),
+            'default_pk_d': bytes(sk.default_pkd()),
+            'v': np.v,
+            'rcm': bytes(np.rcm),
+            'memo': np.memo,
+            'cv': bytes(cv),
+            'cmu': bytes(cm.u),
+            'esk': bytes(esk),
+            'epk': bytes(epk),
+            'shared_secret': bytes(shared_secret),
+            'k_enc': k_enc,
+            'p_enc': p_enc,
+            'c_enc': c_enc,
+            'ock': ock,
+            'op': op,
+            'c_out': c_out,
+        })
+
+    render_tv(
+        args,
+        'sapling_note_encryption',
+        (
+            ('ovk', '[u8; 32]'),
+            ('ivk', '[u8; 32]'),
+            ('default_d', '[u8; 11]'),
+            ('default_pk_d', '[u8; 32]'),
+            ('v', 'u64'),
+            ('rcm', '[u8; 32]'),
+            ('memo', '[u8; 512]'),
+            ('cv', '[u8; 32]'),
+            ('cmu', '[u8; 32]'),
+            ('esk', '[u8; 32]'),
+            ('epk', '[u8; 32]'),
+            ('shared_secret', '[u8; 32]'),
+            ('k_enc', '[u8; 32]'),
+            ('p_enc', '[u8; 564]'),
+            ('c_enc', '[u8; 580]'),
+            ('ock', '[u8; 32]'),
+            ('op', '[u8; 64]'),
+            ('c_out', '[u8; 80]'),
+        ),
+        test_vectors,
+    )
+
+
+if __name__ == '__main__':
+    main()