From ff8eb510c7a5a27c8e576bd86747fffa103fec79 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@electriccoin.co>
Date: Sat, 2 Nov 2024 04:57:07 +0000
Subject: [PATCH] Extract ZIP 32 hardened-only key derivation from Orchard

---
 zcash_test_vectors/orchard/key_components.py | 21 +++++--------
 zcash_test_vectors/zip_0032.py               | 33 ++++++++++++++++++++
 2 files changed, 40 insertions(+), 14 deletions(-)
 create mode 100644 zcash_test_vectors/zip_0032.py

diff --git a/zcash_test_vectors/orchard/key_components.py b/zcash_test_vectors/orchard/key_components.py
index 3a7f864..3d934d3 100755
--- a/zcash_test_vectors/orchard/key_components.py
+++ b/zcash_test_vectors/orchard/key_components.py
@@ -1,10 +1,9 @@
 #!/usr/bin/env python3
 import sys; assert sys.version_info[0] >= 3, "Python 3 required."
 
-from hashlib import blake2b
-
 from ..ff1 import ff1_aes256_encrypt
 from ..sapling.key_components import prf_expand
+from ..zip_0032 import CKDh, HardenedOnlyContext, MKGh
 
 from .generators import NULLIFIER_K_BASE, SPENDING_KEY_BASE, group_hash
 from .pallas import Fp, Scalar, Point
@@ -55,26 +54,20 @@ class SpendingKey(object):
 
 
 class ExtendedSpendingKey(SpendingKey):
+    Orchard = HardenedOnlyContext(b'ZcashIP32Orchard', b'\x81')
+
     def __init__(self, chaincode, data):
         SpendingKey.__init__(self, data)
         self.chaincode = chaincode
 
     @classmethod
     def master(cls, S):
-        digest = blake2b(person=b'ZcashIP32Orchard')
-        digest.update(S)
-        I   = digest.digest()
-        I_L = I[:32]
-        I_R = I[32:]
-        return cls(I_R, I_L)
+        (sk, chaincode) = MKGh(cls.Orchard, S)
+        return cls(chaincode, sk)
 
     def child(self, i):
-        assert 0x80000000 <= i and i <= 0xFFFFFFFF
-
-        I   = prf_expand(self.chaincode, b'\x81' + self.data + i2leosp(32, i))
-        I_L = I[:32]
-        I_R = I[32:]
-        return self.__class__(I_R, I_L)
+        (sk_i, c_i) = CKDh(self.Orchard, self.data, self.chaincode, i)
+        return self.__class__(c_i, sk_i)
 
 
 class FullViewingKey(object):
diff --git a/zcash_test_vectors/zip_0032.py b/zcash_test_vectors/zip_0032.py
new file mode 100644
index 0000000..3466fb0
--- /dev/null
+++ b/zcash_test_vectors/zip_0032.py
@@ -0,0 +1,33 @@
+from hashlib import blake2b
+
+from .sapling.key_components import prf_expand
+from .utils import i2leosp
+
+class HardenedOnlyContext(object):
+    def __init__(self, MKGDomain, CKDDomain):
+        assert type(MKGDomain) == bytes
+        assert len(MKGDomain) == 16
+        assert type(CKDDomain) == bytes
+        assert len(CKDDomain) == 1
+
+        self.MKGDomain = MKGDomain
+        self.CKDDomain = CKDDomain
+
+def MKGh(Context, IKM):
+    assert type(Context) == HardenedOnlyContext
+
+    digest = blake2b(person=Context.MKGDomain)
+    digest.update(IKM)
+    I   = digest.digest()
+    I_L = I[:32]
+    I_R = I[32:]
+    return (I_L, I_R)
+
+def CKDh(Context, sk_par, c_par, i):
+    assert type(Context) == HardenedOnlyContext
+    assert 0x80000000 <= i and i <= 0xFFFFFFFF
+
+    I   = prf_expand(c_par, Context.CKDDomain + sk_par + i2leosp(32, i))
+    I_L = I[:32]
+    I_R = I[32:]
+    return (I_L, I_R)