From 8addd5efff3b22e66d688ca9a4300ed0adf77875 Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Thu, 30 Apr 2026 10:50:45 +0200
Subject: [PATCH] replace the Lioness key derivation by a less a hacky one

---
 reference/Crypto/Lioness.hs   | 18 ++++++------------
 reference/Crypto/Symmetric.hs |  8 ++++++++
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/reference/Crypto/Lioness.hs b/reference/Crypto/Lioness.hs
index 41cddb6..22df5f3 100644
--- a/reference/Crypto/Lioness.hs
+++ b/reference/Crypto/Lioness.hs
@@ -24,7 +24,7 @@ import Octet
 type ByteStream   = [Word8]
 type MasterKey    = Key256
 
-type KeyDerivFun256  = Key256 -> ByteStream
+type KeyDerivFun256  = Domain -> Key256  -> Word256
 type KeyedHashFun256 = Key256 -> [Word8] -> Word256
 type StreamGen256    = Key256 -> ByteStream
 
@@ -44,11 +44,10 @@ type LionessKeys = (Key256,Key256,Key256,Key256)
 
 lionessDeriveKeys :: LionessInstance -> MasterKey -> LionessKeys
 lionessDeriveKeys (MkLioness kdfFun _ _) masterKey = (k1,k2,k3,k4) where
-  [k1,k2,k3,k4] = map (Key256 . W256) $ partition 32 $ take 128 $ kdfFun masterKey
-
-  partition :: Int -> [a] -> [[a]]
-  partition m [] = []
-  partition m xs = take m xs : partition m (drop m xs)
+  k1 = Key256 (kdfFun LionessKey1 masterKey)
+  k2 = Key256 (kdfFun LionessKey2 masterKey)
+  k3 = Key256 (kdfFun LionessKey3 masterKey)
+  k4 = Key256 (kdfFun LionessKey4 masterKey)
 
 --------------------------------------------------------------------------------
 
@@ -109,13 +108,8 @@ lionessInvPerm inst@(MkLioness kdfFun hashFun streamFun) masterKey input
 
 --------------------------------------------------------------------------------
 
-twistIV :: IV -> IV
-twistIV (IV orig) = IV (orig `xor128` twist) where
-  twist = wordFromInteger 0x1234567890abcdef_aa55aa55aa55aa55
-
 testKdfFun :: KeyDerivFun256
-testKdfFun bigKey = case splitKey256 bigKey of 
-  (key,iv) ->  streamCipherPRGBytes AES128_CTR key (twistIV iv)
+testKdfFun domain (Key256 masterKey) = kdf256 KDF_SHA256 domain (fromWord256 masterKey)
 
 testHashFun :: KeyedHashFun256
 testHashFun (Key256 bigKey) input = hash SHA256 (fromWord256 bigKey ++ input)
diff --git a/reference/Crypto/Symmetric.hs b/reference/Crypto/Symmetric.hs
index d00d09e..ca6da40 100644
--- a/reference/Crypto/Symmetric.hs
+++ b/reference/Crypto/Symmetric.hs
@@ -50,6 +50,10 @@ data Domain
   | SphinxMacKey                -- ^ key for the MAC in the Sphinx header
   | SphinxPayloadEncKey         -- ^ key to encrypt the Sphinx payload
   | SphinxBlinding              -- ^ key to compute the blinding factor
+  | LionessKey1                 -- ^ K1 of Lioness
+  | LionessKey2                 -- ^ K2 of Lioness
+  | LionessKey3                 -- ^ K3 of Lioness
+  | LionessKey4                 -- ^ K4 of Lioness
   deriving (Eq,Show)
 
 -- | Key derivation functions
@@ -75,6 +79,10 @@ domainConstant domain =
     SphinxMacKey        -> asciiStringToWord128 "mac-key"
     SphinxPayloadEncKey -> asciiStringToWord128 "payload-enc-key"
     SphinxBlinding      -> asciiStringToWord128 "sphinx-blinding"
+    LionessKey1         -> asciiStringToWord128 "lioness-key1"
+    LionessKey2         -> asciiStringToWord128 "lioness-key2"
+    LionessKey3         -> asciiStringToWord128 "lioness-key3"
+    LionessKey4         -> asciiStringToWord128 "lioness-key4"
   where 
     asciiStringToWord128 :: String -> Word128
     asciiStringToWord128 input