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

-- | Secp256k1 cryptographic primitives.

module Morley.Tezos.Crypto.Secp256k1
  ( -- * Cryptographic primitive types
    PublicKey (..)
  , SecretKey
  , Signature (..)
  , detSecretKey
  , detSecretKeyDo
  , toPublic

  -- * Raw bytes (no checksums, tags or anything)
  , publicKeyToBytes
  , mkPublicKey
  , publicKeyLengthBytes
  , signatureToBytes
  , mkSignature
  , signatureLengthBytes

  -- * Formatting and parsing
  , formatPublicKey
  , mformatPublicKey
  , parsePublicKey
  , formatSignature
  , mformatSignature
  , parseSignature
  , formatSecretKey
  , parseSecretKey

  -- * Signing
  , sign
  , checkSignature
  ) where

import Crypto.Hash (Blake2b_256(..))
import Crypto.PubKey.ECC.ECDSA qualified as ECDSA
import Crypto.PubKey.ECC.Generate qualified as ECC.Generate
import Crypto.PubKey.ECC.Types (Curve(..), CurveName(..), getCurveByName)
import Crypto.PubKey.ECC.Types qualified as ECDSA
import Crypto.Random (MonadRandom)
import Data.ByteArray (ByteArray, ByteArrayAccess)
import Fmt (Buildable, build)

import Morley.Michelson.Text
import Morley.Tezos.Crypto.Util

curve :: Curve
curve = getCurveByName SEC_p256k1

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

-- | Secp256k1 public cryptographic key.
newtype PublicKey = PublicKey
  { unPublicKey :: ECDSA.PublicKey
  } deriving stock (Eq, Show, Generic)

instance NFData PublicKey where
  rnf (PublicKey (ECDSA.PublicKey cu q))
    = rnfCurve cu `seq` rnf q

instance Ord PublicKey where
  compare = compare `on` (publicKeyToBytes :: PublicKey -> ByteString)

-- | Secp256k1 secret cryptographic key.
newtype SecretKey = SecretKey
  { unSecretKey :: ECDSA.KeyPair
  } deriving stock (Show, Eq, Generic)

instance NFData SecretKey where
  rnf (SecretKey (ECDSA.KeyPair cu pp pn)) =
    rnfCurve cu `seq` rnf (pp, pn)

-- | Deterministicaly generate a secret key from seed.
detSecretKey :: ByteString -> SecretKey
detSecretKey seed = deterministic seed $ detSecretKeyDo

detSecretKeyDo :: MonadRandom m => m SecretKey
detSecretKeyDo = SecretKey <$> do
  (publicKey, privateKey) <- ECC.Generate.generate curve
  return $
    ECDSA.KeyPair curve (ECDSA.public_q publicKey) (ECDSA.private_d privateKey)

-- | Create a public key from a secret key.
toPublic :: SecretKey -> PublicKey
toPublic =
  PublicKey .
  ECDSA.PublicKey curve . (\(ECDSA.KeyPair _ pp _) -> pp) . unSecretKey

-- | Secp256k1 cryptographic signature.
newtype Signature = Signature
  { unSignature :: ECDSA.Signature
  } deriving stock (Show, Eq, Generic)

instance NFData Signature where
  rnf (Signature (ECDSA.Signature a b)) = rnf a `seq` rnf b

----------------------------------------------------------------------------
-- Conversion to/from raw bytes (no checksums, tags or anything)
----------------------------------------------------------------------------

-- | Convert a 'PublicKey' to raw bytes.
publicKeyToBytes :: forall ba. ByteArray ba => PublicKey -> ba
publicKeyToBytes (PublicKey p) =
  publicKeyToBytes_ curve p

-- | Make a 'PublicKey' from raw bytes.
mkPublicKey :: ByteArrayAccess ba => ba -> Either CryptoParseError PublicKey
mkPublicKey ba = PublicKey <$> mkPublicKey_ curve ba

publicKeyLengthBytes :: (Integral n, CheckIntSubType Int n) => n
publicKeyLengthBytes = publicKeyLengthBytes_ curve

-- | Convert a 'PublicKey' to raw bytes.
signatureToBytes :: ByteArray ba => Signature -> ba
signatureToBytes (Signature sig) =
  signatureToBytes_ curve sig

-- | Make a 'Signature' from raw bytes.
mkSignature :: ByteArray ba => ba -> Either CryptoParseError Signature
mkSignature ba = Signature <$> mkSignature_ curve ba

signatureLengthBytes :: (Integral n, CheckIntSubType Int n) => n
signatureLengthBytes = signatureLengthBytes_ curve

mkSecretKey :: ByteArray ba => ba -> Either CryptoParseError SecretKey
mkSecretKey = Right . SecretKey . mkSecretKey_ curve

-- | Convert a 'PublicKey' to raw bytes.
secretKeyToBytes :: ByteArray ba => SecretKey -> ba
secretKeyToBytes (SecretKey kp) =
  secretKeyToBytes_ kp

----------------------------------------------------------------------------
-- Magic bytes
----------------------------------------------------------------------------

publicKeyTag :: ByteString
publicKeyTag = "\003\254\226\086"

-- | Base16 format is @11a2e0c9@
secretKeyTag :: ByteString
secretKeyTag = "\017\162\224\201"

signatureTag :: ByteString
signatureTag = "\013\115\101\019\063"

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

formatPublicKey :: PublicKey -> Text
formatPublicKey = formatImpl @ByteString publicKeyTag . publicKeyToBytes

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

instance Buildable PublicKey where
  build = build . formatPublicKey

parsePublicKey :: Text -> Either CryptoParseError PublicKey
parsePublicKey = parseImpl publicKeyTag mkPublicKey

formatSignature :: Signature -> Text
formatSignature = formatImpl @ByteString signatureTag . signatureToBytes

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

instance Buildable Signature where
  build = build . formatSignature

parseSignature :: Text -> Either CryptoParseError Signature
parseSignature = parseImpl signatureTag mkSignature

formatSecretKey :: SecretKey -> Text
formatSecretKey = formatImpl @ByteString secretKeyTag . secretKeyToBytes

instance Buildable SecretKey where
  build = build . formatSecretKey

parseSecretKey :: Text -> Either CryptoParseError SecretKey
parseSecretKey = parseImpl secretKeyTag mkSecretKey

----------------------------------------------------------------------------
-- Signing
----------------------------------------------------------------------------

-- | Sign a message using the secret key.
sign :: MonadRandom m => SecretKey -> ByteString -> m Signature
sign (SecretKey keyPair) msg = fmap Signature $ do
  let pk = ECDSA.toPrivateKey keyPair
  ECDSA.Signature r s' <- ECDSA.sign pk Blake2b_256 msg
  let n = ECDSA.ecc_n . ECDSA.common_curve $ ECDSA.private_curve pk
      -- Tezos uses bitcoin-core implementation of secp256k1, or, more
      -- accurately, nomadic-labs OCaml bindings to bitcoin-core implementation,
      -- see [tezos dependencies] and [C implementation]. It does this thing for
      -- some reason. The signature is valid either way, but the network will
      -- refuse to accept @s > n/2@. note that @s' < n@ due to modular
      -- arithmetic.
      --
      -- [tezos dependencies]: https://gitlab.com/tezos/tezos/-/blob/e9a9b61969b7a44749bed9bd9cdbcb4f2283220f/src/lib_crypto/dune#L16
      -- [C implementation]: https://gitlab.com/nomadic-labs/ocaml-secp256k1-internal/-/blob/1b51aefb3ae579a24529b18cf9da6991b4ff17c2/src/ecdsa_impl.h#L305
      --
      -- Original code for C implementation:
      -- https://github.com/bitcoin/bitcoin/blob/c06cda3e48e9826043ebc5790a7bb505bfbf368c/src/secp256k1/src/ecdsa_impl.h#L305
      s | s' > n `div` 2 = n - s'
        | otherwise = s'
  pure $ ECDSA.Signature r s


-- | Check that a sequence of bytes has been signed with a given key.
checkSignature :: PublicKey -> Signature -> ByteString -> Bool
checkSignature (PublicKey pk) (Signature sig) =
  ECDSA.verify Blake2b_256 pk sig