-- |
-- Module      : Haskoin.Test.Keys
-- Copyright   : No rights reserved
-- License     : MIT
-- Maintainer  : jprupp@protonmail.ch
-- Stability   : experimental
-- Portability : POSIX
module Haskoin.Util.Arbitrary.Keys where

import Crypto.Secp256k1
import Data.Bits (clearBit)
import Data.Coerce (coerce)
import Data.List (foldl')
import Data.Word (Word32)
import Haskoin.Crypto.Hash
import Haskoin.Crypto.Keys.Common
import Haskoin.Crypto.Keys.Extended
import Haskoin.Crypto.Keys.Extended.Internal (Fingerprint (..))
import Haskoin.Crypto.Signature
import Haskoin.Util.Arbitrary.Crypto
import Test.QuickCheck

-- | Arbitrary private key with arbitrary compressed flag.
arbitraryPrivateKey :: Gen PrivateKey
arbitraryPrivateKey :: Gen PrivateKey
arbitraryPrivateKey = Bool -> SecKey -> PrivateKey
wrapSecKey (Bool -> SecKey -> PrivateKey)
-> Gen Bool -> Gen (SecKey -> PrivateKey)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Bool
forall a. Arbitrary a => Gen a
arbitrary Gen (SecKey -> PrivateKey) -> Gen SecKey -> Gen PrivateKey
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen SecKey
forall a. Arbitrary a => Gen a
arbitrary

-- | Arbitrary public key, either compressed or not.
arbitraryPublicKey :: Ctx -> Gen PublicKey
arbitraryPublicKey :: Ctx -> Gen PublicKey
arbitraryPublicKey Ctx
ctx = (PrivateKey, PublicKey) -> PublicKey
forall a b. (a, b) -> b
snd ((PrivateKey, PublicKey) -> PublicKey)
-> Gen (PrivateKey, PublicKey) -> Gen PublicKey
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ctx -> Gen (PrivateKey, PublicKey)
arbitraryKeyPair Ctx
ctx

-- | Arbitrary keypair, both either compressed or not.
arbitraryKeyPair :: Ctx -> Gen (PrivateKey, PublicKey)
arbitraryKeyPair :: Ctx -> Gen (PrivateKey, PublicKey)
arbitraryKeyPair Ctx
ctx = do
  PrivateKey
k <- Gen PrivateKey
arbitraryPrivateKey
  (PrivateKey, PublicKey) -> Gen (PrivateKey, PublicKey)
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return (PrivateKey
k, Ctx -> PrivateKey -> PublicKey
derivePublicKey Ctx
ctx PrivateKey
k)

arbitraryFingerprint :: Gen Fingerprint
arbitraryFingerprint :: Gen Fingerprint
arbitraryFingerprint = Word32 -> Fingerprint
Fingerprint (Word32 -> Fingerprint) -> Gen Word32 -> Gen Fingerprint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word32
forall a. Arbitrary a => Gen a
arbitrary

-- | Arbitrary extended private key.
arbitraryXPrvKey :: Gen XPrvKey
arbitraryXPrvKey :: Gen XPrvKey
arbitraryXPrvKey =
  Word8 -> Fingerprint -> Word32 -> ChainCode -> SecKey -> XPrvKey
XPrvKey
    (Word8 -> Fingerprint -> Word32 -> ChainCode -> SecKey -> XPrvKey)
-> Gen Word8
-> Gen (Fingerprint -> Word32 -> ChainCode -> SecKey -> XPrvKey)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word8
forall a. Arbitrary a => Gen a
arbitrary
    Gen (Fingerprint -> Word32 -> ChainCode -> SecKey -> XPrvKey)
-> Gen Fingerprint
-> Gen (Word32 -> ChainCode -> SecKey -> XPrvKey)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Fingerprint
arbitraryFingerprint
    Gen (Word32 -> ChainCode -> SecKey -> XPrvKey)
-> Gen Word32 -> Gen (ChainCode -> SecKey -> XPrvKey)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Word32
forall a. Arbitrary a => Gen a
arbitrary
    Gen (ChainCode -> SecKey -> XPrvKey)
-> Gen ChainCode -> Gen (SecKey -> XPrvKey)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen ChainCode
arbitraryHash256
    Gen (SecKey -> XPrvKey) -> Gen SecKey -> Gen XPrvKey
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen SecKey
forall a. Arbitrary a => Gen a
arbitrary

-- | Arbitrary extended public key.
arbitraryXPubKey :: Ctx -> Gen XPubKey
arbitraryXPubKey :: Ctx -> Gen XPubKey
arbitraryXPubKey Ctx
ctx = (XPrvKey, XPubKey) -> XPubKey
forall a b. (a, b) -> b
snd ((XPrvKey, XPubKey) -> XPubKey)
-> Gen (XPrvKey, XPubKey) -> Gen XPubKey
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ctx -> Gen (XPrvKey, XPubKey)
arbitraryXKeyPair Ctx
ctx

-- | Arbitrary extended public key with its corresponding private key.
arbitraryXKeyPair :: Ctx -> Gen (XPrvKey, XPubKey)
arbitraryXKeyPair :: Ctx -> Gen (XPrvKey, XPubKey)
arbitraryXKeyPair Ctx
ctx = (\XPrvKey
k -> (XPrvKey
k, Ctx -> XPrvKey -> XPubKey
deriveXPubKey Ctx
ctx XPrvKey
k)) (XPrvKey -> (XPrvKey, XPubKey))
-> Gen XPrvKey -> Gen (XPrvKey, XPubKey)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen XPrvKey
arbitraryXPrvKey

{- Custom derivations -}

-- | Arbitrary derivation index with last bit unset.
genIndex :: Gen Word32
genIndex :: Gen Word32
genIndex = (Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`clearBit` Int
31) (Word32 -> Word32) -> Gen Word32 -> Gen Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word32
forall a. Arbitrary a => Gen a
arbitrary

-- | Arbitrary BIP-32 path index. Can be hardened or not.
arbitraryBip32PathIndex :: Gen Bip32PathIndex
arbitraryBip32PathIndex :: Gen Bip32PathIndex
arbitraryBip32PathIndex =
  [Gen Bip32PathIndex] -> Gen Bip32PathIndex
forall a. [Gen a] -> Gen a
oneof
    [ Word32 -> Bip32PathIndex
Bip32SoftIndex (Word32 -> Bip32PathIndex) -> Gen Word32 -> Gen Bip32PathIndex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word32
genIndex,
      Word32 -> Bip32PathIndex
Bip32HardIndex (Word32 -> Bip32PathIndex) -> Gen Word32 -> Gen Bip32PathIndex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word32
genIndex
    ]

-- | Arbitrary BIP-32 derivation path composed of only hardened derivations.
arbitraryHardPath :: Gen HardPath
arbitraryHardPath :: Gen HardPath
arbitraryHardPath = (HardPath -> Word32 -> HardPath)
-> HardPath -> [Word32] -> HardPath
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' HardPath -> Word32 -> HardPath
forall t. HardOrAny t => DerivPathI t -> Word32 -> DerivPathI t
(:|) HardPath
forall t. DerivPathI t
Deriv ([Word32] -> HardPath) -> Gen [Word32] -> Gen HardPath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word32 -> Gen [Word32]
forall a. Gen a -> Gen [a]
listOf Gen Word32
genIndex

-- | Arbitrary BIP-32 derivation path composed of only non-hardened derivations.
arbitrarySoftPath :: Gen SoftPath
arbitrarySoftPath :: Gen SoftPath
arbitrarySoftPath = (SoftPath -> Word32 -> SoftPath)
-> SoftPath -> [Word32] -> SoftPath
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SoftPath -> Word32 -> SoftPath
forall t. AnyOrSoft t => DerivPathI t -> Word32 -> DerivPathI t
(:/) SoftPath
forall t. DerivPathI t
Deriv ([Word32] -> SoftPath) -> Gen [Word32] -> Gen SoftPath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word32 -> Gen [Word32]
forall a. Gen a -> Gen [a]
listOf Gen Word32
genIndex

-- | Arbitrary derivation path composed of hardened and non-hardened derivations.
arbitraryDerivPath :: Gen DerivPath
arbitraryDerivPath :: Gen DerivPath
arbitraryDerivPath = [Bip32PathIndex] -> DerivPath
concatBip32Segments ([Bip32PathIndex] -> DerivPath)
-> Gen [Bip32PathIndex] -> Gen DerivPath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Bip32PathIndex -> Gen [Bip32PathIndex]
forall a. Gen a -> Gen [a]
listOf Gen Bip32PathIndex
arbitraryBip32PathIndex

-- | Arbitrary parsed derivation path. Can contain 'ParsedPrv', 'ParsedPub' or
-- 'ParsedEmpty' elements.
arbitraryParsedPath :: Gen ParsedPath
arbitraryParsedPath :: Gen ParsedPath
arbitraryParsedPath =
  [Gen ParsedPath] -> Gen ParsedPath
forall a. [Gen a] -> Gen a
oneof
    [ DerivPath -> ParsedPath
ParsedPrv (DerivPath -> ParsedPath) -> Gen DerivPath -> Gen ParsedPath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen DerivPath
arbitraryDerivPath,
      DerivPath -> ParsedPath
ParsedPub (DerivPath -> ParsedPath) -> Gen DerivPath -> Gen ParsedPath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen DerivPath
arbitraryDerivPath,
      DerivPath -> ParsedPath
ParsedEmpty (DerivPath -> ParsedPath) -> Gen DerivPath -> Gen ParsedPath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen DerivPath
arbitraryDerivPath
    ]

-- | Arbitrary message hash, private key, nonce and corresponding signature. The
-- signature is generated with a random message, random private key and a random
-- nonce.
arbitrarySignature :: Ctx -> Gen (Hash256, SecKey, Sig)
arbitrarySignature :: Ctx -> Gen (ChainCode, SecKey, Sig)
arbitrarySignature Ctx
ctx = do
  ChainCode
m <- Gen ChainCode
arbitraryHash256
  SecKey
key <- Gen SecKey
forall a. Arbitrary a => Gen a
arbitrary
  let sig :: Sig
sig = Ctx -> SecKey -> ChainCode -> Sig
signHash Ctx
ctx SecKey
key ChainCode
m
  (ChainCode, SecKey, Sig) -> Gen (ChainCode, SecKey, Sig)
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return (ChainCode
m, SecKey
key, Sig
sig)