Crypto.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Cryptographic primitives used in Tezos.
--
-- WARNING: some functions may be vulnerable to timing attacks.
-- Also, this code was not reviewed by cryptography/security experts.
-- Do not use it with secret keys that have some value.
-- We provide 'SecretKey' type and (limited) signing functionality only
-- for testing.
-- If you need to sign something in production, use something else
-- (e. g. @octez-client@).
--
-- Tezos supports multiple cryptographic curves that are denoted by the
-- number after tz in the public key hash:
-- • tz1 — ed25519
-- • tz2 — secp256k1
-- • tz3 — P256
-- • tz4 — BLS12381
-- We have Morley.Tezos.Crypto.Curve module for each of these curves.
-- They expose very similar functionality and their main purpose is to hide
-- implementation details for each curve as well as some other specifics (e. g.
-- prefixes that are used for human-readable representation).
--
-- Additionally, Tezos uses b2b hashes to represent addresses of contracts
-- (using @KT1@ prefix) and transaction rollups (using @txr1@ prefix) -- these
-- hashes are also implemented here.
--
-- We do not support @txr1@ addresses as those are disabled on the mainnet.
--
-- This module serves two purposes:
-- 1. It is an umbrella module that re-exports some stuff from other modules.
-- 2. Michelson types such as @key@ and @signature@ may store primitive of any
-- curve, so we need "union" types in Haskell as well.
--
-- During conversion to human-readable representation usually some magical
-- prefix is used. They have been found in source code in some repos (e. g.
-- <https://gitlab.com/tezos/tezos/blob/c52ee69231c5ae4d9cec1f3c8aba0c3573922e2a/src/lib_crypto/base58.ml>)
-- and checked manually. Existing tests confirm they are correct.

{-# LANGUAGE DeriveLift #-}

module Morley.Tezos.Crypto
  ( -- * Cryptographic primitive types
    PublicKey (..)
  , SecretKey (..) -- Currently we need to differentiate secret keys in morley-client
  , Signature (..)
  , HashTag (..)
  , KeyHashTag
  , Hash (..)
  , HashKind (..)
  , KeyHash
  , ContractHash
  , SmartRollupHash
  , BLS12381.Bls12381Fr
  , BLS12381.Bls12381G1
  , BLS12381.Bls12381G2

  -- * Public/secret key functions
  , detSecretKey
  , detSecretKey'
  , KeyType(..)
  , keyTypeTag
  , publicKeyType
  , toPublic
  , publicKeyToBytes

  -- * Signature
  , signatureToBytes
  , mkSignature
  , parseSignatureRaw
  , signatureLengthBytes
  , checkSignature
  , sign

  -- * Formatting
  , CryptoParseError (..)
  , formatPublicKey
  , mformatPublicKey
  , parsePublicKey
  , parsePublicKeyRaw
  , formatSignature
  , mformatSignature
  , parseSignature
  , formatHash
  , mformatHash
  , parseHash
  , parseKeyHashRaw
  , hashLengthBytes
  , formatSecretKey
  , parseSecretKey
  , decodeKeyHash

  -- * Hashing
  , hashKey
  , blake2b
  , blake2b160
  , keccak
  , sha256
  , sha3
  , sha512

  -- * Timelock puzzle
  , Chest
  , ChestKey
  , OpeningResult(..)
  , TLTime(..)
  , openChest
  , mkTLTime
  , toTLTime

  -- * Utilities
  , encodeBase58Check
  , decodeBase58Check
  , B58CheckWithPrefixError (..)
  , decodeBase58CheckWithPrefix
  , parseSomeHashBase58
  , keyDecoders
  , keyHashDecoders
  , AllHashTags(..)
  ) where

import Control.Monad.Except (throwError)
import Crypto.Number.Serialize (os2ip)
import Crypto.Random (MonadRandom)
import Data.Aeson (FromJSON(..), FromJSONKey, ToJSON(..), ToJSONKey)
import Data.Aeson qualified as Aeson
import Data.Aeson.Encoding qualified as Aeson
import Data.Aeson.Types qualified as AesonTypes
import Data.Binary.Get qualified as Get
import Data.ByteArray qualified as BA
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Data.Some (Some(..), mapSome)
import Data.Text qualified as T
import Fmt (Buildable, build, hexF, pretty)
import Instances.TH.Lift ()
import Language.Haskell.TH.Syntax (Lift)

import Morley.Michelson.Text
import Morley.Tezos.Crypto.BLS qualified as BLS
import Morley.Tezos.Crypto.BLS12381 qualified as BLS12381
import Morley.Tezos.Crypto.Ed25519 qualified as Ed25519
import Morley.Tezos.Crypto.Hash
import Morley.Tezos.Crypto.P256 qualified as P256
import Morley.Tezos.Crypto.Secp256k1 qualified as Secp256k1
import Morley.Tezos.Crypto.Timelock
  (Chest, ChestKey, OpeningResult(..), TLTime(..), mkTLTime, openChest, toTLTime)
import Morley.Tezos.Crypto.Util
import Morley.Util.Binary
import Morley.Util.CLI
import Morley.Util.TH (deriveGADTNFData)

----------------------------------------------------------------------------
-- Types, instances, conversions
----------------------------------------------------------------------------

-- | A kind of a hash.
data HashKind
  = HashKindPublicKey -- ^ Public key hash for @tz1@, @tz2@, @tz3@ addresses.
  | HashKindContract -- ^ Contract hash for @KT1@ smart contract addresses.
  | HashKindSmartRollup -- ^ Smart rollup hash for @sr1@ addresses.

-- | Type of public/secret key as enum.
data KeyType
  = KeyTypeEd25519
  | KeyTypeSecp256k1
  | KeyTypeP256
  | KeyTypeBLS
  deriving stock (Show, Eq, Enum, Bounded, Ord, Lift, Generic)
  deriving anyclass NFData

instance Buildable KeyType where
  build = \case
    KeyTypeEd25519   -> "key Ed25519"
    KeyTypeSecp256k1 -> "key Secp256k1"
    KeyTypeP256      -> "key P256"
    KeyTypeBLS       -> "key BLS"

-- | What specific type of hash is used for the 'Hash'.
data HashTag (kind :: HashKind) where
  HashKey :: KeyType -> HashTag 'HashKindPublicKey
  HashContract :: HashTag 'HashKindContract
  HashSR :: HashTag 'HashKindSmartRollup

deriving stock instance Show (HashTag kind)
deriving stock instance Eq (HashTag kind)
deriving stock instance Ord (HashTag kind)
deriving stock instance Lift (HashTag kind)

-- NB: these definitions are here and not below because TH scoping rules are annoying

deriveGADTNFData ''HashTag

-- | Public cryptographic key used by Tezos.
-- There are three cryptographic curves each represented by its own constructor.
data PublicKey
  = PublicKeyEd25519 Ed25519.PublicKey
  -- ^ Public key that uses the ed25519 cryptographic curve.
  | PublicKeySecp256k1 Secp256k1.PublicKey
  -- ^ Public key that uses the secp256k1 cryptographic curve.
  | PublicKeyP256 P256.PublicKey
  -- ^ Public key that uses the NIST P-256 cryptographic curve.
  | PublicKeyBLS BLS.PublicKey
  -- ^ Public key that uses the BLS12-381 cryptographic curve.
  deriving stock (Show, Eq, Ord, Generic)

instance NFData PublicKey

-- | Secret cryptographic key used by Tezos.
-- Constructors correspond to 'PublicKey' constructors.
data SecretKey
  = SecretKeyEd25519 Ed25519.SecretKey
  -- ^ Secret key that uses the ed25519 cryptographic curve.
  | SecretKeySecp256k1 Secp256k1.SecretKey
  -- ^ Secret key that uses the secp256k1 cryptographic curve.
  | SecretKeyP256 P256.SecretKey
  -- ^ Secret key that uses the NIST P-256 cryptographic curve.
  | SecretKeyBLS  BLS.SecretKey
  -- ^ Secret key that uses BLS12-381 curve.
  deriving stock (Show, Eq, Generic)

instance NFData SecretKey

instance HasCLReader SecretKey where
  getReader = eitherReader (first pretty . parseSecretKey . toText)
  getMetavar = "SECRET_KEY"

-- | Deterministically generate a secret key from seed. Key type is specified
-- explicitly.
detSecretKey' :: KeyType -> ByteString -> SecretKey
detSecretKey' = \case
  KeyTypeEd25519   -> SecretKeyEd25519   . Ed25519.detSecretKey
  KeyTypeSecp256k1 -> SecretKeySecp256k1 . Secp256k1.detSecretKey
  KeyTypeP256      -> SecretKeyP256      . P256.detSecretKey
  KeyTypeBLS       -> SecretKeyBLS       . BLS.detSecretKey

-- | Deterministically generate a secret key from seed. Type of the key depends
-- on seed value.
detSecretKey :: ByteString -> SecretKey
detSecretKey seed = detSecretKey'
  (toEnum $ fromIntegralOverflowing (os2ip seed) `mod` (fromEnum (maxBound :: KeyType) + 1))
  seed

-- | Create a public key from a secret key.
toPublic :: SecretKey -> PublicKey
toPublic = \case
  SecretKeyEd25519 sk -> PublicKeyEd25519 . Ed25519.toPublic $ sk
  SecretKeySecp256k1 sk -> PublicKeySecp256k1 . Secp256k1.toPublic $ sk
  SecretKeyP256 sk -> PublicKeyP256 . P256.toPublic $ sk
  SecretKeyBLS sk -> PublicKeyBLS . BLS.toPublic $ sk

-- | Cryptographic signatures used by Tezos.
-- Constructors correspond to 'PublicKey' constructors.
--
-- Tezos distinguishes signatures for different curves.
-- For instance, ed25519 signatures and secp256k1 signatures
-- are printed differently (have different prefix).
-- However, signatures are packed without information about the
-- curve. For this purpose there is a generic signature which
-- only stores bytes and doesn't carry information about the curve.
-- Apparently unpacking from bytes always produces such signature.
-- Unpacking from string produces a signature with curve information.
data Signature
  = SignatureEd25519 Ed25519.Signature
  -- ^ Signature that uses the ed25519 cryptographic curve.
  | SignatureSecp256k1 Secp256k1.Signature
  -- ^ Siganture that uses the secp256k1 cryptographic curve.
  | SignatureP256 P256.Signature
  -- ^ Signature that uses the NIST P-256 cryptographic curve.
  | SignatureBLS BLS.Signature
  -- ^ Signature that uses the BLS12-381 cryptographic curve.
  | SignatureGeneric ByteString
  -- ^ Generic signature for which curve is unknown.
  deriving stock (Show, Generic)

instance NFData Signature

-- This instance slightly differs from the default one. If one
-- signature is generic and the other one is not, they still may be
-- equal if they have the same byte representation.
-- With default instance packing a signature and unpacking it would produce
-- a different (with respect to 'Eq') signature which is inconvenient.
instance Eq Signature where
  sig1 == sig2 = case (sig1, sig2) of
    (SignatureGeneric bytes1, _) -> bytes1 == signatureToBytes sig2
    (_, SignatureGeneric bytes2) -> signatureToBytes sig1 == bytes2

    (SignatureEd25519 s1, SignatureEd25519 s2) -> s1 == s2
    (SignatureEd25519 {}, _) -> False

    (SignatureSecp256k1 s1, SignatureSecp256k1 s2) -> s1 == s2
    (SignatureSecp256k1 {}, _) -> False

    (SignatureP256 s1, SignatureP256 s2) -> s1 == s2
    (SignatureP256 {}, _) -> False

    (SignatureBLS s1, SignatureBLS s2) -> s1 == s2
    (SignatureBLS {}, _) -> False

instance Ord Signature where
  compare = compare `on` signatureToBytes @ByteString

----------------------------------------------------------------------------
-- Signature
----------------------------------------------------------------------------

-- | Convert a 'Signature' to raw bytes.
signatureToBytes :: BA.ByteArray ba => Signature -> ba
signatureToBytes = \case
  SignatureEd25519 sig -> Ed25519.signatureToBytes sig
  SignatureSecp256k1 sig -> Secp256k1.signatureToBytes sig
  SignatureP256 sig -> P256.signatureToBytes sig
  SignatureBLS sig -> BLS.signatureToBytes sig
  SignatureGeneric bytes -> BA.convert bytes

-- | Make a 'Signature' from raw bytes.
-- Can return only 'SignatureGeneric' or 'SignatureBLS'
mkSignature :: BA.ByteArray ba => ba -> Maybe Signature
mkSignature ba
  | l == signatureLengthBytes = Just $ SignatureGeneric $ BA.convert ba
  | l == BLS.signatureLengthBytes = SignatureBLS <$> rightToMaybe (BLS.mkSignature ba)
  | otherwise = Nothing
  where
   l = BA.length ba

parseSignatureRaw :: ByteString -> Either ParseSignatureRawError Signature
parseSignatureRaw ba = maybeToRight (ParseSignatureRawWrongSize ba) $ mkSignature ba

data ParseSignatureRawError
  = ParseSignatureRawWrongSize ByteString
  deriving stock (Eq, Show, Generic)

instance Buildable ParseSignatureRawError where
  build =
    \case
      ParseSignatureRawWrongSize ba -> "Given raw signature " <>
        hexF ba <> " has invalid length " <> show (length ba)

-- Apparently Tezos relies on the fact that in all schemes signature
-- size is 64 bytes, so it also has generic signature and always reads
-- 64 bytes during unpack.
-- So we can have one 'signatureLengthBytes' and do not have to
-- distinguish between curves.
-- However, we still have such a check here just in case as a precaution.
signatureLengthBytes :: HasCallStack => Integral n => n
signatureLengthBytes
  | all is64
    [ Ed25519.signatureLengthBytes
    , P256.signatureLengthBytes
    , Secp256k1.signatureLengthBytes
    ]
  = 64
  | otherwise =
    error "Apparently our understanding of signatures in Tezos is broken"
  where
    is64 :: Int -> Bool
    is64 = (== 64)

genericSignatureTag :: ByteString
genericSignatureTag = "\004\130\043"

-- | Check that a sequence of bytes has been signed with a given key.
checkSignature :: PublicKey -> Signature -> ByteString -> Bool
checkSignature pk0 sig0 bytes =
  case (pk0, sig0) of
    (PublicKeyEd25519 pk, SignatureEd25519 sig) ->
      Ed25519.checkSignature pk sig bytes
    (PublicKeySecp256k1 pk, SignatureSecp256k1 sig) ->
      Secp256k1.checkSignature pk sig bytes
    (PublicKeyP256 pk, SignatureP256 sig) ->
      P256.checkSignature pk sig bytes
    (PublicKeyBLS pk, SignatureBLS sig) ->
      BLS.checkSignature pk sig bytes
    (PublicKeyEd25519 pk, SignatureGeneric sBytes) ->
      case Ed25519.mkSignature sBytes of
        Right sig -> Ed25519.checkSignature pk sig bytes
        Left _ -> False
    (PublicKeySecp256k1 pk, SignatureGeneric sBytes) ->
      case Secp256k1.mkSignature sBytes of
        Right sig -> Secp256k1.checkSignature pk sig bytes
        Left _ -> False
    (PublicKeyP256 pk, SignatureGeneric sBytes) ->
      case P256.mkSignature sBytes of
        Right sig -> P256.checkSignature pk sig bytes
        Left _ -> False
    _ -> False

sign :: MonadRandom m => SecretKey -> ByteString -> m Signature
sign sk bs =
  case sk of
    SecretKeyEd25519 sk'   -> pure $ SignatureEd25519 $ Ed25519.sign sk' bs
    SecretKeySecp256k1 sk' -> SignatureSecp256k1 <$> Secp256k1.sign sk' bs
    SecretKeyP256 sk'      -> SignatureP256 <$> P256.sign sk' bs
    SecretKeyBLS sk'       -> pure $ SignatureBLS $ BLS.sign sk' bs

----------------------------------------------------------------------------
-- Formatting
----------------------------------------------------------------------------

formatPublicKey :: PublicKey -> Text
formatPublicKey = \case
  PublicKeyEd25519 pk -> Ed25519.formatPublicKey pk
  PublicKeySecp256k1 pk -> Secp256k1.formatPublicKey pk
  PublicKeyP256 pk -> P256.formatPublicKey pk
  PublicKeyBLS pk -> BLS.formatPublicKey pk

mformatPublicKey :: PublicKey -> MText
mformatPublicKey = unsafe . mkMText . formatPublicKey

instance Buildable PublicKey where
  build = build . formatPublicKey

parsePublicKey :: Text -> Either CryptoParseError PublicKey
parsePublicKey txt =
  firstRight $ map ($ txt)
    ( fmap PublicKeyEd25519 . Ed25519.parsePublicKey :|
    [ fmap PublicKeySecp256k1 . Secp256k1.parsePublicKey
    , fmap PublicKeyP256 . P256.parsePublicKey
    , fmap PublicKeyBLS . BLS.parsePublicKey
    ])

parsePublicKeyRaw :: ByteString -> Either Text PublicKey
parsePublicKeyRaw ba = bimap (toText . view _3) (view _3) $
  Get.runGetOrFail (decodeWithTag "key" keyDecoders) (LBS.fromStrict ba)

formatSignature :: Signature -> Text
formatSignature = \case
  SignatureEd25519 sig -> Ed25519.formatSignature sig
  SignatureSecp256k1 sig -> Secp256k1.formatSignature sig
  SignatureP256 sig -> P256.formatSignature sig
  SignatureBLS sig -> BLS.formatSignature sig
  SignatureGeneric sig -> formatImpl genericSignatureTag sig

mformatSignature :: Signature -> MText
mformatSignature = unsafe . mkMText . formatSignature

instance Buildable Signature where
  build = build . formatSignature

parseSignature :: Text -> Either CryptoParseError Signature
parseSignature txt =
  firstRight $ map ($ txt)
    ( fmap SignatureEd25519 . Ed25519.parseSignature :|
    [ fmap SignatureSecp256k1 . Secp256k1.parseSignature
    , fmap SignatureP256 . P256.parseSignature
    , fmap SignatureBLS . BLS.parseSignature
    , parseImpl genericSignatureTag (pure . SignatureGeneric)
    ])

formatSecretKey :: SecretKey -> Text
formatSecretKey key = "unencrypted:" <> case key of
  SecretKeyEd25519 sig -> Ed25519.formatSecretKey sig
  SecretKeySecp256k1 sig -> Secp256k1.formatSecretKey sig
  SecretKeyP256 sig -> P256.formatSecretKey sig
  SecretKeyBLS sig -> BLS.formatSecretKey sig

instance Buildable SecretKey where
  build = build . formatSecretKey

-- | Parse __unencrypted__ secret key. It accepts formats containing
-- either with or without the @unecrypted@ prefix.
parseSecretKey :: Text -> Either CryptoParseError SecretKey
parseSecretKey txt =
  firstRight $ map (\f -> f $ removePrefix txt)
    ( fmap SecretKeyEd25519 . Ed25519.parseSecretKey :|
    [ fmap SecretKeySecp256k1 . Secp256k1.parseSecretKey
    , fmap SecretKeyP256 . P256.parseSecretKey
    , fmap SecretKeyBLS . BLS.parseSecretKey
    ])
  where
    removePrefix :: Text -> Text
    removePrefix input =
      let unencrypted = "unencrypted:"
          (prefix, payload) = T.splitAt (length unencrypted) input
      in case prefix == unencrypted of
        True -> payload
        False -> input

----------------------------------------------------------------------------
-- JSON encoding/decoding
----------------------------------------------------------------------------

-- If you ever need these instances for any particular 'PublicKey' or
-- 'Signature', you can define them in respective modules the same
-- way.

instance ToJSON PublicKey where
  toJSON = Aeson.String . formatPublicKey
  toEncoding = Aeson.text . formatPublicKey

instance FromJSON PublicKey where
  parseJSON =
    Aeson.withText "PublicKey" $
    either (fail . pretty) pure . parsePublicKey

instance ToJSON Signature where
  toJSON = Aeson.String . formatSignature
  toEncoding = Aeson.text . formatSignature

instance FromJSON Signature where
  parseJSON =
    Aeson.withText "Signature" $
    either (fail . pretty) pure . parseSignature

instance ToJSON (Hash kind) where
  toJSON = Aeson.String . formatHash
  toEncoding = Aeson.text . formatHash

instance ToJSONKey (Hash kind) where
  toJSONKey = AesonTypes.toJSONKeyText formatHash

instance AllHashTags kind => FromJSON (Hash kind) where
  parseJSON =
    Aeson.withText "Hash" $
    either (fail . pretty) pure . parseHash

instance AllHashTags kind => FromJSONKey (Hash kind) where
  fromJSONKey =
    AesonTypes.FromJSONKeyTextParser $
    either (fail . pretty) pure . parseHash

----------------------------------------------------------------------------
-- Hash
----------------------------------------------------------------------------

-- | A compatibility synonym for a public key hash tag.
type KeyHashTag = HashTag 'HashKindPublicKey

-- | List all 'HashTag's for a given 'HashKind'.
class AllHashTags kind where
  allHashTags :: NonEmpty (HashTag kind)

instance AllHashTags 'HashKindPublicKey where
  allHashTags = HashKey <$> minBound :| [succ minBound .. maxBound]

instance AllHashTags 'HashKindContract where
  allHashTags = pure HashContract

instance AllHashTags 'HashKindSmartRollup where
  allHashTags = pure HashSR

-- | Blake2b_160 hash of something.
data Hash (kind :: HashKind) = Hash
  { hTag :: HashTag kind
  -- ^ Which kind of hash.
  , hBytes :: ByteString
  -- ^ Hash itself.
  } deriving stock (Show, Eq, Ord, Generic, Lift)

instance NFData (Hash kind)

-- | Convenience synonym for an on-chain public key hash.
type KeyHash = Hash 'HashKindPublicKey

-- | Convenience synonym for a contract hash.
type ContractHash = Hash 'HashKindContract

-- | Convenience synonym for a smart rollup hash.
type SmartRollupHash = Hash 'HashKindSmartRollup

-- | Length of a hash in bytes (only the hash itself, no tags, checksums
-- or anything).
hashLengthBytes :: Integral n => n
hashLengthBytes = 20

-- | Compute the b58check of a public key hash.
hashKey :: PublicKey -> KeyHash
hashKey pk = Hash (HashKey (publicKeyType pk)) $ blake2b160 $ publicKeyToBytes pk

publicKeyToBytes :: PublicKey -> ByteString
publicKeyToBytes = \case
  PublicKeyEd25519 pk -> Ed25519.publicKeyToBytes pk
  PublicKeySecp256k1 pk -> Secp256k1.publicKeyToBytes pk
  PublicKeyP256 pk -> P256.publicKeyToBytes pk
  PublicKeyBLS pk -> BLS.publicKeyToBytes pk

mkPublicKey :: BA.ByteArray ba => KeyType -> ba -> Either CryptoParseError PublicKey
mkPublicKey = \case
  KeyTypeEd25519   -> fmap PublicKeyEd25519 . Ed25519.mkPublicKey
  KeyTypeSecp256k1 -> fmap PublicKeySecp256k1 . Secp256k1.mkPublicKey
  KeyTypeP256      -> fmap PublicKeyP256 . P256.mkPublicKey
  KeyTypeBLS       -> fmap PublicKeyBLS . BLS.mkPublicKey

formatHash :: (Hash kind) -> Text
formatHash (Hash tag bytes) = formatImpl (hashTagBytes tag) bytes

mformatHash :: (Hash kind) -> MText
mformatHash = unsafe . mkMText . formatHash

instance Buildable (Hash kind) where
  build = build . formatHash

parseHash
  :: AllHashTags kind
  => Text
  -> Either CryptoParseError (Hash kind)
parseHash txt =
  let
    mkKeyHash tag bs =
      Hash tag bs <$
        unless (length bs == hashLengthBytes)
        (Left $ CryptoParseUnexpectedLength "KeyHash" (length bs))

    parse :: HashTag kind -> Either CryptoParseError (Hash kind)
    parse tag = mkKeyHash tag =<< parseImpl (hashTagBytes tag) pure txt

  in firstRight $ map parse allHashTags

-- | Parse a 'Hash' of any known kind from its its human-readable textual representation.
parseSomeHashBase58 :: Text -> Either CryptoParseError (Some Hash)
parseSomeHashBase58 = maybe (Left CryptoParseWrongBase58Check) parseSomeHash .  decodeBase58Check

parseSomeHash :: ByteString -> Either CryptoParseError (Some Hash)
parseSomeHash a = do
  (sometag, bs) <- parseSomeHashTag a
  when (BS.length bs /= hashLengthBytes) $
    Left (CryptoParseUnexpectedLength "KeyHash" (length bs))
  return $ mapSome (flip Hash bs) sometag

parseKeyHashHelper
  :: Int
  -> LText
  -> ExceptT CryptoParseError Get.Get a
  -> ByteString
  -> Either CryptoParseError a
parseKeyHashHelper expectedLength name decoder ba
  | BS.length ba /= expectedLength
  = Left $ CryptoParseUnexpectedLength name (BS.length ba)
  | otherwise
  = either (Left . CryptoParseBinaryError . toText . view _3) (view _3)
  $ flip Get.runGetOrFail (LBS.fromStrict ba) $ runExceptT decoder

parseKeyHashRaw :: ByteString -> Either CryptoParseError KeyHash
parseKeyHashRaw = parseKeyHashHelper (hashLengthBytes + 1) "key_hash" decodeKeyHash

-- | Magic constants used by Tezos to encode hashes with proper prefixes.
hashTagBytes :: HashTag kind -> ByteString
hashTagBytes =
  \case
    HashKey KeyTypeEd25519 -> "\006\161\159" -- tz1
    -- https://gitlab.com/tezos/tezos/-/blob/0ca82c9dc361a6f223e81221c86bdb95d1a8d91c/src/lib_crypto/base58.ml#L379
    HashKey KeyTypeSecp256k1 -> "\006\161\161" -- tz2
    -- https://gitlab.com/tezos/tezos/-/blob/0ca82c9dc361a6f223e81221c86bdb95d1a8d91c/src/lib_crypto/base58.ml#L381
    HashKey KeyTypeP256 -> "\006\161\164" -- tz3
    -- https://gitlab.com/tezos/tezos/-/blob/0ca82c9dc361a6f223e81221c86bdb95d1a8d91c/src/lib_crypto/base58.ml#L383
    HashKey KeyTypeBLS -> "\006\161\166" -- tz4
    -- https://gitlab.com/tezos/tezos/-/blob/4b0dd9e9715ce82ac6429571d8843ab681522daf/src/lib_crypto/base58.ml#L371
    HashContract -> "\2\90\121" -- KT1
    -- https://gitlab.com/tezos/tezos/-/blob/0ca82c9dc361a6f223e81221c86bdb95d1a8d91c/src/proto_alpha/lib_protocol/contract_hash.ml#L27
    HashSR -> "\006\124\117" -- sr1
    -- https://gitlab.com/tezos/tezos/-/blob/f7f6047237974ef85d94c87368f2a82615bcc8ca/src/proto_016_PtMumbai/lib_protocol/sc_rollup_repr.ml#L33

parseSomeHashTag :: ByteString -> Either CryptoParseError (Some HashTag, ByteString)
parseSomeHashTag bs = maybeToRight failHash $ asum
  [ tryHash (HashKey KeyTypeEd25519)
  , tryHash (HashKey KeyTypeSecp256k1)
  , tryHash (HashKey KeyTypeP256)
  , tryHash (HashKey KeyTypeBLS)
  , tryHash HashContract
  , tryHash HashSR
  ]
  where
    tryHash :: HashTag kind -> Maybe (Some HashTag, ByteString)
    tryHash hashKind = (Some hashKind,) <$> BS.stripPrefix (hashTagBytes hashKind) bs

    hashTXR = "\001\128\120\031" -- txr1
    -- https://gitlab.com/tezos/tezos/-/blob/0ca82c9dc361a6f223e81221c86bdb95d1a8d91c/src/proto_014_PtKathma/lib_protocol/tx_rollup_prefixes.ml#L35

    failHash :: CryptoParseError
    failHash
      | hashTXR `BS.isPrefixOf` bs = CryptoParseUnsupportedTag "txr1" bs
      | otherwise = CryptoParseWrongTag bs

instance AllHashTags kind => HasCLReader (Hash kind) where
  getReader = eitherReader (first pretty . parseHash . toText)
  getMetavar = "KEY_HASH"

keyTypeTag :: KeyType -> Word8
keyTypeTag = \case
  KeyTypeEd25519   -> 0x00
  KeyTypeSecp256k1 -> 0x01
  KeyTypeP256      -> 0x02
  KeyTypeBLS       -> 0x03

publicKeyType :: PublicKey -> KeyType
publicKeyType = \case
  PublicKeyEd25519{}   -> KeyTypeEd25519
  PublicKeySecp256k1{} -> KeyTypeSecp256k1
  PublicKeyP256{}      -> KeyTypeP256
  PublicKeyBLS{}       -> KeyTypeBLS

keyDecoders :: [TaggedDecoder PublicKey]
keyDecoders = [minBound..] <&> mkKeyParser
  where
    mkKeyParser x = keyTypeTag x #: decodeBytesLike (pretty x) (mkPublicKey x)

keyHashDecoders :: (Monad (t Get.Get), MonadTrans t) => [TaggedDecoderM t KeyHash]
keyHashDecoders = [minBound..] <&> mkKeyHashParser
  where
    mkKeyHashParser kt = keyTypeTag kt ##: Hash (HashKey kt) <$> getPayload
    getPayload = lift $ getByteStringCopy hashLengthBytes

decodeKeyHash :: ExceptT CryptoParseError Get.Get KeyHash
decodeKeyHash =
  decodeWithTagM "key_hash" (throwError . CryptoParseWrongTag . BS.singleton) keyHashDecoders