-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | 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. `tezos-client`).
--
-- Tezos supports 3 cryptographic curves that are denoted by the
-- number after tz in the public key hash: tz1, tz2 or tz3.
-- • tz1 — ed25519
-- • tz2 — secp256k1
-- • tz3 — P256
-- We have 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).
--
-- 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.

module Tezos.Crypto
  ( -- * Cryptographic primitive types
    PublicKey (..)
  , SecretKey (..) -- Currently we need to differentiate secret keys in morley-client
  , Signature (..)
  , KeyHashTag (..)
  , KeyHash (..)

  -- * Public/secret key functions
  , detSecretKey
  , toPublic

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

  -- * Formatting
  , CryptoParseError (..)
  , formatPublicKey
  , mformatPublicKey
  , parsePublicKey
  , parsePublicKeyRaw
  , formatSignature
  , mformatSignature
  , parseSignature
  , formatKeyHash
  , mformatKeyHash
  , parseKeyHash
  , parseKeyHashRaw
  , keyHashLengthBytes
  , formatSecretKey
  , parseSecretKey

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

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

import Crypto.Random (MonadRandom)
import Data.Aeson (FromJSON(..), ToJSON(..))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Encoding as Aeson
import qualified Data.Binary.Get as Get
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Text as T
import Fmt (Buildable, build, hexF, pretty)

import Michelson.Text
import qualified Tezos.Crypto.Ed25519 as Ed25519
import Tezos.Crypto.Hash
import qualified Tezos.Crypto.P256 as P256
import qualified Tezos.Crypto.Secp256k1 as Secp256k1
import Tezos.Crypto.Util
import Util.Binary
import Util.CLI

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

-- | 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.
  deriving stock (Show, Eq, 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.
  deriving stock (Show, Eq, Generic)

instance NFData SecretKey

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

-- | Deterministicaly generate a secret key from seed.
-- Type of the key depends on seed length.
detSecretKey :: HasCallStack => ByteString -> SecretKey
detSecretKey seed = seed & case (length seed + 2) `mod` 3 of
  0 -> SecretKeyEd25519 . Ed25519.detSecretKey
  1 -> SecretKeySecp256k1 . Secp256k1.detSecretKey
  2 -> SecretKeyP256 . P256.detSecretKey
  _ -> error "detSecretKey: unexpected happened"

-- | 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

-- | 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.
  | 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, SignatureGeneric bytes2) -> bytes1 == bytes2
    (SignatureGeneric bytes1, SignatureEd25519 (Ed25519.signatureToBytes -> bytes2)) ->
      bytes1 == bytes2
    (SignatureGeneric bytes1, SignatureSecp256k1 (Secp256k1.signatureToBytes -> bytes2)) ->
      bytes1 == bytes2
    (SignatureGeneric bytes1, SignatureP256 (P256.signatureToBytes -> bytes2)) ->
      bytes1 == bytes2

    (_, SignatureGeneric {}) -> sig2 == sig1

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

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

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

----------------------------------------------------------------------------
-- 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
  SignatureGeneric bytes -> BA.convert bytes

-- | Make a 'Signature' from raw bytes.
-- Can return only generic signature.
mkSignature :: BA.ByteArray ba => ba -> Maybe Signature
mkSignature ba =
  SignatureGeneric (BA.convert ba) <$ guard (l == signatureLengthBytes)
  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
    (PublicKeyEd25519 pk, SignatureGeneric sBytes) ->
      Ed25519.mkSignature sBytes
        & fmap (\sig -> Ed25519.checkSignature pk sig bytes)
        & fromRight False
    (PublicKeySecp256k1 pk, SignatureGeneric sBytes) ->
      Secp256k1.mkSignature sBytes
        & fmap (\sig -> Secp256k1.checkSignature pk sig bytes)
        & fromRight False
    (PublicKeyP256 pk, SignatureGeneric sBytes) ->
      P256.mkSignature sBytes
        & fmap (\sig -> P256.checkSignature pk sig bytes)
        & fromRight 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

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

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

mformatPublicKey :: PublicKey -> MText
mformatPublicKey = mkMTextUnsafe . 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
    ])

parsePublicKeyRaw :: ByteString -> Either Text PublicKey
parsePublicKeyRaw ba =
  case Get.runGetOrFail (decodeWithTag "key" keyDecoders) (LBS.fromStrict ba) of
    Right (_, _, result) -> Right result
    Left (_, _, err) -> Left (toText err)

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

mformatSignature :: Signature -> MText
mformatSignature = mkMTextUnsafe . 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
    , 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

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
    ])
  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 KeyHash where
  toJSON = Aeson.String . formatKeyHash
  toEncoding = Aeson.text . formatKeyHash

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

----------------------------------------------------------------------------
-- KeyHash
----------------------------------------------------------------------------

-- | Which curve was used for the hashed public key inside 'KeyHash'.
data KeyHashTag
  = KeyHashEd25519
  | KeyHashSecp256k1
  | KeyHashP256
  deriving stock (Show, Eq, Ord, Bounded, Enum, Generic)

instance NFData KeyHashTag

-- | Blake2b_160 hash of a public key.
data KeyHash = KeyHash
  { khTag :: KeyHashTag
  -- ^ We store which curve was used because it affects formatting.
  , khBytes :: ByteString
  -- ^ Hash itself.
  } deriving stock (Show, Eq, Ord, Generic)

instance NFData KeyHash

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

-- | Compute the b58check of a public key hash.
hashKey :: PublicKey -> KeyHash
hashKey =
  \case
    PublicKeyEd25519 pk ->
      KeyHash KeyHashEd25519 (blake2b160 $ Ed25519.publicKeyToBytes pk)
    PublicKeySecp256k1 pk ->
      KeyHash KeyHashSecp256k1 (blake2b160 $ Secp256k1.publicKeyToBytes pk)
    PublicKeyP256 pk ->
      KeyHash KeyHashP256 (blake2b160 $ P256.publicKeyToBytes pk)

formatKeyHash :: KeyHash -> Text
formatKeyHash (KeyHash tag bytes) = formatImpl (keyHashTagBytes tag) bytes

mformatKeyHash :: KeyHash -> MText
mformatKeyHash = mkMTextUnsafe . formatKeyHash

instance Buildable KeyHash where
  build = build . formatKeyHash

parseKeyHash :: Text -> Either CryptoParseError KeyHash
parseKeyHash txt =
  let
    mkKeyHash tag bs =
      KeyHash tag bs <$
        unless (length bs == keyHashLengthBytes)
        (Left $ CryptoParseUnexpectedLength "KeyHash" (length bs))

    parse :: KeyHashTag -> Either CryptoParseError KeyHash
    parse tag = mkKeyHash tag =<< parseImpl (keyHashTagBytes tag) pure txt

  in firstRight $ map parse $ minBound :| [succ minBound ..]

parseKeyHashRaw :: ByteString -> Either CryptoParseError KeyHash
parseKeyHashRaw ba =
  if (BS.length ba - 1 == keyHashLengthBytes) then
    case Get.runGetOrFail (decodeWithTag "key_hash" keyHashDecoders) (LBS.fromStrict ba) of
      Right (_, _, result) -> Right result
      Left (_, _, err) -> Left $ CryptoParseBinaryError (toText err)
  else
    Left $ CryptoParseUnexpectedLength "key_hash" (BS.length ba)

keyHashTagBytes :: KeyHashTag -> ByteString
keyHashTagBytes =
  \case
    KeyHashEd25519 -> "\006\161\159"
    KeyHashSecp256k1 -> "\006\161\161"
    KeyHashP256 -> "\006\161\164"

instance HasCLReader KeyHash where
  getReader = eitherReader (first pretty . parseKeyHash . toText)
  getMetavar = "KEY_HASH"

keyDecoders :: [TaggedDecoder PublicKey]
keyDecoders =
  [ 0x00 #: decodeBytesLike "key Ed25519"
      (fmap PublicKeyEd25519 . Ed25519.mkPublicKey)
  , 0x01 #: decodeBytesLike "key Secp256k1"
      (fmap PublicKeySecp256k1 . Secp256k1.mkPublicKey)
  , 0x02 #: decodeBytesLike "key P256"
      (fmap PublicKeyP256 . P256.mkPublicKey)
  ]

keyHashDecoders :: [TaggedDecoder KeyHash]
keyHashDecoders =
  [ 0x00 #: KeyHash KeyHashEd25519 <$> getPayload
  , 0x01 #: KeyHash KeyHashSecp256k1 <$> getPayload
  , 0x02 #: KeyHash KeyHashP256 <$> getPayload
  ]
  where
    getPayload = getByteStringCopy keyHashLengthBytes