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

-- | Utilities shared by multiple cryptographic primitives.

module Tezos.Crypto.Util
  ( CryptoParseError (..)
  , encodeBase58Check
  , decodeBase58Check
  , B58CheckWithPrefixError (..)
  , decodeBase58CheckWithPrefix
  , formatImpl
  , parseImpl
  , firstRight
  , deterministic

  -- * ECDSA Utils
  , rnfCurve
  , publicKeyLengthBytes_
  , mkSignature_
  , mkSecretKey_
  , secretKeyToBytes_
  , signatureToBytes_
  , mkPublicKey_
  , publicKeyToBytes_
  , signatureLengthBytes_
  ) where

import Crypto.Error (CryptoError)
import Crypto.Number.ModArithmetic (squareRoot)
import Crypto.Number.Serialize (i2ospOf_, os2ip)
import qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import qualified Crypto.PubKey.ECC.Generate as ECC.Generate
import Crypto.PubKey.ECC.Types
  (Curve(..), CurveCommon(..), CurvePrime(..), Point(..), curveSizeBits)
import Crypto.Random (ChaChaDRG, MonadPseudoRandom, drgNewSeed, seedFromInteger, withDRG)
import qualified Data.Binary.Get as Get
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base58 as Base58
import qualified Data.ByteString.Lazy as LBS
import Fmt (Buildable, Builder, build, fmt, hexF, (+|), (|+))

import Tezos.Crypto.Hash
import Util.Binary (getRemainingByteStringCopy)

-- | Error that can happen during parsing of cryptographic primitive types.
data CryptoParseError
  = CryptoParseWrongBase58Check
  | CryptoParseWrongTag ByteString
  | CryptoParseCryptoError CryptoError
  | CryptoParseUnexpectedLength Builder Int
  | CryptoParseBinaryError Text
  deriving stock (Show, Eq)

instance NFData CryptoParseError where
  rnf = rnf @String . show

instance Buildable CryptoParseError where
  build =
    \case
      CryptoParseWrongBase58Check -> "Wrong base58check encoding of bytes"
      CryptoParseWrongTag tag -> "Prefix is wrong tag: " <> fmt (hexF tag)
      CryptoParseCryptoError err ->
        "Cryptographic library reported an error: " <>
        build (displayException err)
      CryptoParseUnexpectedLength what l ->
        "Unexpected length of " <> what <> ": " <> build l
      CryptoParseBinaryError err -> "" +| err |+ ""

-- | Encode a bytestring in Base58Check format.
encodeBase58Check :: ByteString -> Text
encodeBase58Check =
  decodeUtf8 . Base58.encodeBase58 Base58.bitcoinAlphabet . withCheckSum
  where
    withCheckSum :: ByteString -> ByteString
    withCheckSum bs = bs <> checkSum bs

-- | Decode a bytestring from Base58Check format.
decodeBase58Check :: Text -> Maybe ByteString
decodeBase58Check text = do
  bytes <- Base58.decodeBase58 Base58.bitcoinAlphabet (encodeUtf8 text)
  let (payload, chk) = BS.splitAt (length bytes - 4) bytes
  guard $ chk == checkSum payload
  return payload

checkSum :: ByteString -> ByteString
checkSum = BS.take 4 . (sha256 . sha256)

data B58CheckWithPrefixError
  = B58CheckWithPrefixWrongPrefix ByteString
  | B58CheckWithPrefixWrongEncoding
  deriving stock (Show)

-- | Parse a base58check encoded value expecting some prefix. If the
-- actual prefix matches the expected one, it's stripped of and the
-- resulting payload is returned.
decodeBase58CheckWithPrefix ::
  ByteString -> Text -> Either B58CheckWithPrefixError ByteString
decodeBase58CheckWithPrefix prefix text =
  case decodeBase58Check text of
    Nothing -> Left B58CheckWithPrefixWrongEncoding
    Just bs ->
      let (actualPrefix, payload) = BS.splitAt (length prefix) bs
       in if actualPrefix == prefix
          then Right payload
          else Left (B58CheckWithPrefixWrongPrefix actualPrefix)

-- | Template for 'format*' functions.
formatImpl :: BA.ByteArrayAccess x => ByteString -> x -> Text
formatImpl tag = encodeBase58Check . mappend tag . BA.convert

-- | Template for 'parse*' functions.
parseImpl
  :: ByteString
  -> (ByteString -> Either CryptoParseError res)
  -> Text
  -> Either CryptoParseError res
parseImpl expectedTag constructor text = do
  let convertErr :: B58CheckWithPrefixError -> CryptoParseError
      convertErr =
        \case B58CheckWithPrefixWrongPrefix prefix -> CryptoParseWrongTag prefix
              B58CheckWithPrefixWrongEncoding -> CryptoParseWrongBase58Check
  payload <- first convertErr $ decodeBase58CheckWithPrefix expectedTag text
  constructor payload

-- | Returns first encountered 'Right' in a list. If there are none,
-- returns arbitrary 'Left'.
-- It is useful to implement parsing.
firstRight :: NonEmpty (Either e a) -> Either e a
firstRight (h :| rest) =
  case h of
    Left e -> maybe (Left e) firstRight $ nonEmpty rest
    Right a -> Right a

-- | Do randomized action using specified seed.
deterministic :: ByteString -> MonadPseudoRandom ChaChaDRG a -> a
deterministic seed = fst . withDRG chachaSeed
  where
    chachaSeed = drgNewSeed . seedFromInteger . os2ip $ seed

---------------------------------------------------------
-- Utilities shared by @Secp256k1@ and @P256@.
---------------------------------------------------------

rnfCurve :: Curve -> ()
rnfCurve cu =
  case cu of
    CurveF2m c -> rnf c
    CurveFP (CurvePrime i (CurveCommon a b c d e)) ->
      rnf (i, a, b, c, d, e)

curveSizeBytes :: Curve -> Int
curveSizeBytes curve = curveSizeBits curve `div` 8

signatureLengthBytes_ :: Integral n => Curve -> n
signatureLengthBytes_ curve = fromIntegral $ (curveSizeBytes curve) + (curveSizeBytes curve)

coordToBytes :: BA.ByteArray ba => Curve -> Integer -> ba
coordToBytes curve = i2ospOf_ (curveSizeBytes curve)

publicKeyLengthBytes_ :: Integral n => Curve -> n
publicKeyLengthBytes_ curve = fromIntegral $ 1 + (curveSizeBytes curve)

-- | Make a 'PublicKey' from raw bytes.
--
-- Raw bytes are in the format of Compressed SEC Format. Refer to this article on how this is parsed:
-- <https://www.oreilly.com/library/view/programming-bitcoin/9781492031482/ch04.html>
--
mkPublicKey_ :: BA.ByteArrayAccess ba => Curve -> ba -> Either CryptoParseError ECDSA.PublicKey
mkPublicKey_ curve ba
  | l == (publicKeyLengthBytes_ curve) = do
      (isYEven, x) <- toCryptoEither $ Get.runGetOrFail getX
          (LBS.fromStrict $ BA.convert ba)
      (p, a, b) <- fromCurveFP curve
      let alpha = x ^ (3 :: Integer) + a * x + b
      beta <- squareRoot p alpha
                & maybeToRight (CryptoParseBinaryError "Could not find square root.")
      let (evenBeta, oddBeta) =
            if even beta then
              (beta, p - beta)
            else
              (p - beta, beta)
      let y = if isYEven then evenBeta
                         else oddBeta
      pure $ ECDSA.PublicKey curve $ Point x y
  | otherwise =
    Left $ CryptoParseUnexpectedLength "public key" l
  where
    l = BA.length ba

    getX :: Get.Get (Bool, Integer)
    getX = do
      yPrefix <- Get.getWord8
      xBytes <- getRemainingByteStringCopy
      return (even yPrefix, os2ip xBytes)

    fromCurveFP :: Curve -> Either CryptoParseError (Integer, Integer, Integer)
    fromCurveFP = \case
      CurveFP (CurvePrime p (CurveCommon a b _ _ _)) -> Right (p, a, b)
      CurveF2m _ -> Left $ CryptoParseBinaryError
        "Should not happen. Expect `curve` to be `CurveFP` but got `CurveF2m` instead."

    toCryptoEither :: Either (_a, _b, String) (_c, _d, a) -> Either CryptoParseError a
    toCryptoEither g =
      case g of
        Right (_, _, a) -> Right a
        Left (_, _, err) -> Left $ CryptoParseBinaryError $ toText err

-- | Convert a 'PublicKey' to raw bytes.
publicKeyToBytes_ :: forall ba. (BA.ByteArray ba, HasCallStack) => Curve -> ECDSA.PublicKey -> ba
publicKeyToBytes_ curve (ECDSA.PublicKey _ publicPoint) =
  case publicPoint of
    Point x y -> prefix y `BA.append` coordToBytes curve x
    PointO -> error "PublicKey somehow contains infinity point"
  where
    prefix :: Integer -> ba
    prefix y
      | odd y = BA.singleton 0x03
      | otherwise = BA.singleton 0x02

-- | Convert a 'PublicKey' to raw bytes.
signatureToBytes_ :: BA.ByteArray ba => Curve -> ECDSA.Signature -> ba
signatureToBytes_ curve (ECDSA.Signature r s) =
  coordToBytes curve r <> coordToBytes curve s

-- | Convert a 'PublicKey' to raw bytes.
secretKeyToBytes_ :: BA.ByteArray ba => ECDSA.KeyPair -> ba
secretKeyToBytes_ (ECDSA.KeyPair c _ s) =
  coordToBytes c s

-- | Make a 'Signature' from raw bytes.
mkSignature_ :: BA.ByteArray ba => Curve -> ba -> Either CryptoParseError ECDSA.Signature
mkSignature_ curve ba
  | l == (signatureLengthBytes_ curve)
  , (rBytes, sBytes) <- BA.splitAt (curveSizeBytes curve) ba =
    Right $ ECDSA.Signature (os2ip rBytes) (os2ip sBytes)
  | otherwise =
    Left $ CryptoParseUnexpectedLength "signature" l
  where
    l = BA.length ba

-- | Make a 'SecretKey' from raw bytes.
mkSecretKey_ :: BA.ByteArray ba => Curve -> ba -> ECDSA.KeyPair
mkSecretKey_ c ba =
  let s = os2ip ba
      p = ECC.Generate.generateQ c s
  in ECDSA.KeyPair c p s