{-|
Module: Crypto.Spake2.Util
Description: Miscellany. Mostly to do with serialization.
-}
module Crypto.Spake2.Util
  ( expandData
  , expandArbitraryElementSeed
  , bytesToNumber
  , numberToBytes
  , unsafeNumberToBytes
  ) where

import Protolude

import Crypto.Hash.Algorithms (SHA256)
import Crypto.Number.Serialize (os2ip, i2ospOf, i2ospOf_)
import qualified Crypto.KDF.HKDF as HKDF
import Data.ByteArray (ByteArray, ByteArrayAccess(..))

-- | Take an arbitrary sequence of bytes and expand it to be the given number
-- of bytes. Do this by extracting a pseudo-random key and expanding it using
-- HKDF.
expandData :: (ByteArrayAccess input, ByteArray output) => ByteString -> input -> Int -> output
expandData :: ByteString -> input -> Int -> output
expandData ByteString
info input
input Int
size =
  PRK SHA256 -> ByteString -> Int -> output
forall a info out.
(HashAlgorithm a, ByteArrayAccess info, ByteArray out) =>
PRK a -> info -> Int -> out
HKDF.expand PRK SHA256
prk ByteString
info Int
size
  where
    prk :: HKDF.PRK SHA256
    prk :: PRK SHA256
prk = ByteString -> input -> PRK SHA256
forall a salt ikm.
(HashAlgorithm a, ByteArrayAccess salt, ByteArrayAccess ikm) =>
salt -> ikm -> PRK a
HKDF.extract ByteString
salt input
input

    -- XXX: I'm no crypto expert, but hard-coding an empty string as a salt
    -- seems kind of weird.
    salt :: ByteString
    salt :: ByteString
salt = ByteString
""

-- | Given a seed value for an arbitrary element (see 'arbitraryElement'),
-- expand it to be of the given length.
expandArbitraryElementSeed :: (ByteArrayAccess ikm, ByteArray out) => ikm -> Int -> out
expandArbitraryElementSeed :: ikm -> Int -> out
expandArbitraryElementSeed =
  -- NOTE: This must be exactly this string in order to interoperate with python-spake2
  ByteString -> ikm -> Int -> out
forall input output.
(ByteArrayAccess input, ByteArray output) =>
ByteString -> input -> Int -> output
expandData ByteString
"SPAKE2 arbitrary element"


-- | Serialize a number according to the SPAKE2 protocol.
--
-- Just kidding, there isn't a SPAKE2 protocol.
-- This just matches the Python implementation.
--
-- Inverse of 'bytesToNumber'.
numberToBytes :: ByteArray bytes => Int -> Integer -> Maybe bytes
numberToBytes :: Int -> Integer -> Maybe bytes
numberToBytes = Int -> Integer -> Maybe bytes
forall ba. ByteArray ba => Int -> Integer -> Maybe ba
i2ospOf

-- | Serialize a number according to the SPAKE2 protocol.
--
-- Panics if the number is too big to fit into the given number of bytes.
unsafeNumberToBytes :: ByteArray bytes => Int -> Integer -> bytes
unsafeNumberToBytes :: Int -> Integer -> bytes
unsafeNumberToBytes = Int -> Integer -> bytes
forall ba. ByteArray ba => Int -> Integer -> ba
i2ospOf_


-- | Deserialize a number according to the SPAKE2 protocol.
--
-- Just kidding, there isn't a SPAKE2 protocol.
-- This just matches the Python implementation.
--
-- Inverse of 'numberToBytes'.
bytesToNumber :: ByteArrayAccess bytes => bytes -> Integer
bytesToNumber :: bytes -> Integer
bytesToNumber = bytes -> Integer
forall ba. ByteArrayAccess ba => ba -> Integer
os2ip