module Tendermint.SDK.Crypto
  ( MakeDigest(..)
  , SignatureSchema(..)
  , RecoverableSignatureSchema(..)
  , parsePubKey
  , Secp256k1
  ) where

import           Control.Error                          (note)
import           Crypto.Hash                            (Digest, hashWith)
import           Crypto.Hash.Algorithms                 (Keccak_256 (..),
                                                         SHA256)
import qualified Crypto.Secp256k1                       as Secp256k1
import           Data.ByteArray                         (convert)
import qualified Data.ByteArray.Base64String            as Base64
import qualified Data.ByteString                        as B
import qualified Data.ByteString.Short                  as Short
import           Data.Kind                              (Type)
import           Data.Maybe                             (fromMaybe)
import           Data.Proxy
import           Data.Text                              (Text)
import qualified Network.ABCI.Types.Messages.FieldTypes as FT
import           Tendermint.SDK.Types.Address           (Address,
                                                         addressFromBytes)

-- | Class encapsulating data which can hashed.
class MakeDigest a where
  makeDigest :: a -> Digest SHA256

-- | Defines the types and methods for the signature schema parameterized by 'alg'.
class SignatureSchema alg where
    type PubKey alg :: Type
    type PrivateKey alg :: Type
    type Signature alg :: Type
    type Message alg :: Type

    algorithm :: Proxy alg -> Text
    sign :: Proxy alg -> PrivateKey alg -> Message alg -> Signature alg
    verify :: Proxy alg -> PubKey alg -> Signature alg -> Message alg -> Bool

    makePubKey :: Proxy alg -> B.ByteString -> Maybe (PubKey alg)
    makeSignature :: Proxy alg -> B.ByteString -> Maybe (Signature alg)
    derivePubKey :: Proxy alg -> PrivateKey alg -> PubKey alg
    addressFromPubKey :: Proxy alg -> PubKey alg -> Address

-- | Class allowing for signing and recovering signatures for messages.
class SignatureSchema alg => RecoverableSignatureSchema alg where
    type RecoverableSignature alg :: Type

    signRecoverableMessage :: Proxy alg -> PrivateKey alg -> Message alg -> RecoverableSignature alg
    recover :: Proxy alg -> RecoverableSignature alg -> Message alg -> Maybe (PubKey alg)
    serializeRecoverableSignature :: Proxy alg -> RecoverableSignature alg -> B.ByteString
    makeRecoverableSignature :: Proxy alg -> B.ByteString -> Maybe (RecoverableSignature alg)

data Secp256k1

msgFromSHA256 :: Digest SHA256 -> Secp256k1.Msg
msgFromSHA256 :: Digest SHA256 -> Msg
msgFromSHA256 dig :: Digest SHA256
dig = Msg -> Maybe Msg -> Msg
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Msg
forall a. HasCallStack => [Char] -> a
error "Digest SHA256 wasn't 32 bytes.") (Maybe Msg -> Msg) -> Maybe Msg -> Msg
forall a b. (a -> b) -> a -> b
$
  ByteString -> Maybe Msg
Secp256k1.msg (ByteString -> Maybe Msg) -> ByteString -> Maybe Msg
forall a b. (a -> b) -> a -> b
$ Digest SHA256 -> ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert Digest SHA256
dig

instance SignatureSchema Secp256k1 where
    type PubKey Secp256k1 = Secp256k1.PubKey
    type PrivateKey Secp256k1 = Secp256k1.SecKey
    type Signature Secp256k1 = Secp256k1.Sig
    type Message Secp256k1 = Digest SHA256

    algorithm :: Proxy Secp256k1 -> Text
algorithm _ = "secp256k1"
    sign :: Proxy Secp256k1
-> PrivateKey Secp256k1 -> Message Secp256k1 -> Signature Secp256k1
sign _ priv :: PrivateKey Secp256k1
priv dig :: Message Secp256k1
dig = SecKey -> Msg -> Sig
Secp256k1.signMsg SecKey
PrivateKey Secp256k1
priv (Digest SHA256 -> Msg
msgFromSHA256 Digest SHA256
Message Secp256k1
dig)
    verify :: Proxy Secp256k1
-> PubKey Secp256k1
-> Signature Secp256k1
-> Message Secp256k1
-> Bool
verify _ pub :: PubKey Secp256k1
pub sig :: Signature Secp256k1
sig dig :: Message Secp256k1
dig = PubKey -> Sig -> Msg -> Bool
Secp256k1.verifySig PubKey
PubKey Secp256k1
pub Sig
Signature Secp256k1
sig (Digest SHA256 -> Msg
msgFromSHA256 Digest SHA256
Message Secp256k1
dig)

    makePubKey :: Proxy Secp256k1 -> ByteString -> Maybe (PubKey Secp256k1)
makePubKey _ = ByteString -> Maybe PubKey
ByteString -> Maybe (PubKey Secp256k1)
Secp256k1.importPubKey
    makeSignature :: Proxy Secp256k1 -> ByteString -> Maybe (Signature Secp256k1)
makeSignature _ = ByteString -> Maybe Sig
ByteString -> Maybe (Signature Secp256k1)
Secp256k1.importSig
    -- For lack of a better idea, we're just going to use the Ethereum style here
    derivePubKey :: Proxy Secp256k1 -> PrivateKey Secp256k1 -> PubKey Secp256k1
derivePubKey _ = SecKey -> PubKey
PrivateKey Secp256k1 -> PubKey Secp256k1
Secp256k1.derivePubKey
    addressFromPubKey :: Proxy Secp256k1 -> PubKey Secp256k1 -> Address
addressFromPubKey _ = ByteString -> Address
addressFromBytes (ByteString -> Address)
-> (PubKey -> ByteString) -> PubKey -> Address
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
B.drop 12 (ByteString -> ByteString)
-> (PubKey -> ByteString) -> PubKey -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Digest Keccak_256 -> ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert (Digest Keccak_256 -> ByteString)
-> (PubKey -> Digest Keccak_256) -> PubKey -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      Keccak_256 -> ByteString -> Digest Keccak_256
forall ba alg.
(ByteArrayAccess ba, HashAlgorithm alg) =>
alg -> ba -> Digest alg
hashWith Keccak_256
Keccak_256 (ByteString -> Digest Keccak_256)
-> (PubKey -> ByteString) -> PubKey -> Digest Keccak_256
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PubKey -> ByteString
Secp256k1.exportPubKey Bool
False

instance RecoverableSignatureSchema Secp256k1 where
    type RecoverableSignature Secp256k1 = Secp256k1.RecSig

    signRecoverableMessage :: Proxy Secp256k1
-> PrivateKey Secp256k1
-> Message Secp256k1
-> RecoverableSignature Secp256k1
signRecoverableMessage _ priv :: PrivateKey Secp256k1
priv dig :: Message Secp256k1
dig = SecKey -> Msg -> RecSig
Secp256k1.signRecMsg SecKey
PrivateKey Secp256k1
priv (Digest SHA256 -> Msg
msgFromSHA256 Digest SHA256
Message Secp256k1
dig)
    recover :: Proxy Secp256k1
-> RecoverableSignature Secp256k1
-> Message Secp256k1
-> Maybe (PubKey Secp256k1)
recover _ sig :: RecoverableSignature Secp256k1
sig dig :: Message Secp256k1
dig = RecSig -> Msg -> Maybe PubKey
Secp256k1.recover RecSig
RecoverableSignature Secp256k1
sig (Digest SHA256 -> Msg
msgFromSHA256 Digest SHA256
Message Secp256k1
dig)
    serializeRecoverableSignature :: Proxy Secp256k1 -> RecoverableSignature Secp256k1 -> ByteString
serializeRecoverableSignature _ sig :: RecoverableSignature Secp256k1
sig =
      let csr :: CompactRecSig
csr = RecSig -> CompactRecSig
Secp256k1.exportCompactRecSig RecSig
RecoverableSignature Secp256k1
sig
      in ShortByteString -> ByteString
Short.fromShort (CompactRecSig -> ShortByteString
Secp256k1.getCompactRecSigR CompactRecSig
csr) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<>
           ShortByteString -> ByteString
Short.fromShort (CompactRecSig -> ShortByteString
Secp256k1.getCompactRecSigS CompactRecSig
csr) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<>
           [Word8] -> ByteString
B.pack [CompactRecSig -> Word8
Secp256k1.getCompactRecSigV CompactRecSig
csr]
    makeRecoverableSignature :: Proxy Secp256k1
-> ByteString -> Maybe (RecoverableSignature Secp256k1)
makeRecoverableSignature _ bs :: ByteString
bs =
      let (r :: ByteString
r,rest :: ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt 32 ByteString
bs
          (s :: ByteString
s,v :: ByteString
v) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt 32 ByteString
rest
      in if ByteString -> Int
B.length ByteString
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= 32 Bool -> Bool -> Bool
|| ByteString -> Int
B.length ByteString
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= 32 Bool -> Bool -> Bool
|| ByteString -> Int
B.length ByteString
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= 1
           then Maybe (RecoverableSignature Secp256k1)
forall a. Maybe a
Nothing
           else CompactRecSig -> Maybe RecSig
CompactRecSig -> Maybe (RecoverableSignature Secp256k1)
Secp256k1.importCompactRecSig (CompactRecSig -> Maybe (RecoverableSignature Secp256k1))
-> CompactRecSig -> Maybe (RecoverableSignature Secp256k1)
forall a b. (a -> b) -> a -> b
$
                  ShortByteString -> ShortByteString -> Word8 -> CompactRecSig
Secp256k1.CompactRecSig (ByteString -> ShortByteString
Short.toShort ByteString
r) (ByteString -> ShortByteString
Short.toShort ByteString
s) (ByteString -> Word8
B.head ByteString
v)


parsePubKey
  :: SignatureSchema alg
  => Proxy alg
  -> FT.PubKey
  -> Either Text (PubKey alg)
parsePubKey :: Proxy alg -> PubKey -> Either Text (PubKey alg)
parsePubKey p :: Proxy alg
p FT.PubKey{..}
  | Text
pubKeyType Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Proxy alg -> Text
forall k (alg :: k). SignatureSchema alg => Proxy alg -> Text
algorithm Proxy alg
p =
      Text -> Maybe (PubKey alg) -> Either Text (PubKey alg)
forall a b. a -> Maybe b -> Either a b
note "Couldn't parse PubKey" (Maybe (PubKey alg) -> Either Text (PubKey alg))
-> Maybe (PubKey alg) -> Either Text (PubKey alg)
forall a b. (a -> b) -> a -> b
$ Proxy alg -> ByteString -> Maybe (PubKey alg)
forall k (alg :: k).
SignatureSchema alg =>
Proxy alg -> ByteString -> Maybe (PubKey alg)
makePubKey Proxy alg
p (Base64String -> ByteString
forall ba. ByteArray ba => Base64String -> ba
Base64.toBytes Base64String
pubKeyData)
  | Bool
otherwise = Text -> Either Text (PubKey alg)
forall a b. a -> Either a b
Left (Text -> Either Text (PubKey alg))
-> Text -> Either Text (PubKey alg)
forall a b. (a -> b) -> a -> b
$ "Unsupported curve: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
pubKeyType