-- | 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
  , Signature (..)
  , KeyHashTag (..)
  , KeyHash (..)

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

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

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

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

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

import Data.Aeson (FromJSON(..), ToJSON(..))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Encoding as Aeson
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import Fmt (Buildable, build, pretty)
import Test.QuickCheck (Arbitrary(..), elements, oneof)

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

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

instance Arbitrary PublicKey where
  arbitrary = toPublic <$> arbitrary

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

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

instance Arbitrary SecretKey where
  arbitrary = oneof
    [ SecretKeyEd25519 <$> arbitrary
    , SecretKeySecp256k1 <$> arbitrary
    , SecretKeyP256 <$> arbitrary
    ]

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

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

instance Arbitrary Signature where
  arbitrary = oneof
    [ SignatureEd25519 <$> arbitrary
    , SignatureSecp256k1 <$> arbitrary
    , SignatureP256 <$> arbitrary
    , SignatureGeneric . BS.replicate signatureLengthBytes <$> arbitrary
    ]

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

-- 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.
-- TODO (#18) consider generic signature here as well.
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
    _ -> False

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

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

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

instance Arbitrary KeyHashTag where
  arbitrary = elements [minBound .. ]

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

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

instance Arbitrary KeyHash where
  arbitrary = hashKey <$> arbitrary

-- | 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
    parse :: KeyHashTag -> Either CryptoParseError KeyHash
    parse tag = KeyHash tag <$> parseImpl (keyHashTagBytes tag) pure txt

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

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