Util.hs

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

-- | Utilities shared by multiple cryptographic primitives.

module Morley.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 Debug qualified (show)

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

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


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

instance NFData CryptoParseError where
  rnf = rnf @String . Debug.show

instance Buildable CryptoParseError where
  build = \case
    CryptoParseWrongBase58Check -> "Wrong base58check encoding of bytes"
    CryptoParseWrongTag tag -> "Prefix is wrong tag:" <+> hexF tag
    CryptoParseUnsupportedTag name tag ->
      "Unsupported hash" <+> build name <> ":" <+> hexF tag
    CryptoParseCryptoError err ->
      "Cryptographic library reported an error:" <+> build (displayException err)
    CryptoParseBLSTError err ->
      "Cryptographic library reported an error:" <+> build (displayException err)
    CryptoParseUnexpectedLength what l ->
      "Unexpected length of" <+> build what <> ":" <+> build l
    CryptoParseBinaryError err -> build 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 base58text = do
  bytes <- Base58.decodeBase58 Base58.bitcoinAlphabet (encodeUtf8 base58text)
  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 base58text =
  case decodeBase58Check base58text 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 textToParse = do
  let convertErr :: B58CheckWithPrefixError -> CryptoParseError
      convertErr =
        \case B58CheckWithPrefixWrongPrefix prefix -> CryptoParseWrongTag prefix
              B58CheckWithPrefixWrongEncoding -> CryptoParseWrongBase58Check
  payload <- first convertErr $ decodeBase58CheckWithPrefix expectedTag textToParse
  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, CheckIntSubType Int n) => Curve -> n
signatureLengthBytes_ curve = fromIntegral $ 2 * (curveSizeBytes curve)

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

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

-- | Make a 'ECDSA.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 'ECDSA.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 'ECDSA.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 'ECDSA.PublicKey' to raw bytes.
secretKeyToBytes_ :: BA.ByteArray ba => ECDSA.KeyPair -> ba
secretKeyToBytes_ (ECDSA.KeyPair c _ s) =
  coordToBytes c s

-- | Make a 'ECDSA.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 'ECDSA.KeyPair' from raw bytes representing a secret key.
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