{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-|
Module      : Crypto.Secp256k1
Description : Public SECP256K1 cryptographic functions
License     : PublicDomain
Maintainer  : root@haskoin.com
Stability   : experimental
Portability : POSIX

This module exposes crytpographic functions from Bitcoin’s secp256k1 library.
Depends on <https://github.com/bitcoin/secp256k1 secp256k1>.
-}
module Crypto.Secp256k1
( -- * Messages
  Msg, msg, getMsg
  -- * Secret Key
, SecKey, importSecKey, exportSecKey, pubKey
  -- ** Raw Secret Key
, secKey, getSecKey
  -- * Public Key
, PubKey, importPubKey, exportPubKey
  -- * Signature
, Sig, CompactSig(..)
, importSig, laxImportSig, exportSig
, exportCompactSig, importCompactSig
, signMsg, verifySig, normalizeSig
  -- * Addition & Multiplication
, Tweak, tweak, getTweak
, tweakAddSecKey, tweakMulSecKey
, tweakAddPubKey, tweakMulPubKey
, combinePubKeys
) where

import           Control.Applicative
import           Control.Monad
import           Crypto.Secp256k1.Internal
import           Data.ByteString           (ByteString)
import qualified Data.ByteString           as BS
import qualified Data.ByteString.Base16    as B16
import           Data.Maybe
import           Data.String
import           Data.String.Conversions
import           Foreign
import           System.IO.Unsafe
import           Test.QuickCheck
import           Text.Read

-- | Internal public key data type.
newtype PubKey = PubKey (ForeignPtr PubKey64)

-- | Internal message data type.
newtype Msg = Msg (ForeignPtr Msg32)

-- | Internal signature data type.
newtype Sig = Sig (ForeignPtr Sig64)

-- | Internal secret key data type.
newtype SecKey = SecKey (ForeignPtr SecKey32)

-- | Internal tweak data type for addition and multiplication.
newtype Tweak = Tweak (ForeignPtr Tweak32)

decodeHex :: ConvertibleStrings a ByteString => a -> Maybe ByteString
decodeHex str = if BS.null r then Just bs else Nothing where
    (bs, r) = B16.decode $ cs str

-- TODO: Test
instance Read PubKey where
    readPrec = parens $ do
        Ident "PubKey" <- lexP
        String str <- lexP
        maybe pfail return $ importPubKey =<< decodeHex str

-- TODO: Test
instance IsString PubKey where
    fromString = fromJust . (importPubKey <=< decodeHex)

-- TODO: Test
instance Show PubKey where
    showsPrec d k = showParen (d > 10) $
        showString "PubKey " . shows (B16.encode $ exportPubKey True k)

-- TODO: Test
instance Read Msg where
    readPrec = parens $ do
        Ident "Msg" <- lexP
        String str <- lexP
        maybe pfail return $ msg =<< decodeHex str

-- TODO: Test
instance IsString Msg where
    fromString = fromJust . msg . cs

-- TODO: Test
instance Show Msg where
    showsPrec d m = showParen (d > 10) $
        showString "Msg " . shows (B16.encode $ getMsg m)

-- TODO: Test
instance Read Sig where
    readPrec = parens $ do
        Ident "Sig" <- lexP
        String str <- lexP
        maybe pfail return $ importSig =<< decodeHex str

-- TODO: Test
instance IsString Sig where
    fromString = fromJust . (importSig <=< decodeHex)

-- TODO: Test
instance Show Sig where
    showsPrec d s = showParen (d > 10) $
        showString "Sig " . shows (B16.encode $ exportSig s)

-- TODO: Test
instance Read SecKey where
    readPrec = parens $ do
        Ident "SecKey" <- lexP
        String str <- lexP
        maybe pfail return $ secKey =<< decodeHex str

-- TODO: Test
instance IsString SecKey where
    fromString str = fromJust $
        (secKey =<< decodeHex str) <|> (importSecKey =<< decodeHex str)

-- TODO: Test
instance Show SecKey where
    showsPrec d k = showParen (d > 10) $
        showString "SecKey " . shows (B16.encode $ getSecKey k)

-- TODO: Test
instance Read Tweak where
    readPrec = parens $ do
        Ident "Tweak" <- lexP
        String str <- lexP
        maybe pfail return $ tweak =<< decodeHex str

-- TODO: Test
instance IsString Tweak where
    fromString = fromJust . (tweak <=< decodeHex)

-- TODO: Test
instance Show Tweak where
    showsPrec d k = showParen (d > 10) $
        showString "Tweak " . shows (B16.encode $ getTweak k)

instance Eq PubKey where
    fp1 == fp2 = getPubKey fp1 == getPubKey fp2

instance Eq Msg where
    fm1 == fm2 = getMsg fm1 == getMsg fm2

instance Eq Sig where
    fg1 == fg2 = exportCompactSig fg1 == exportCompactSig fg2

instance Eq SecKey where
    fk1 == fk2 = getSecKey fk1 == getSecKey fk2

instance Eq Tweak where
    ft1 == ft2 = getTweak ft1 == getTweak ft2

-- | Create internal message data from 32-byte 'ByteString'.
msg :: ByteString -> Maybe Msg
msg bs
    | BS.length bs == 32 = unsafePerformIO $ do
        fp <- mallocForeignPtr
        withForeignPtr fp $ flip poke (Msg32 bs)
        return $ Just $ Msg fp
    | otherwise = Nothing

-- | Create internal secret key data from 32-byte 'ByteString'.
secKey :: ByteString -> Maybe SecKey
secKey bs
    | BS.length bs == 32 = unsafePerformIO $ do
        fp <- mallocForeignPtr
        ret <- withForeignPtr fp $ \p -> do
            poke p (SecKey32 bs)
            ec_seckey_verify ctx p
        if isSuccess ret
            then return $ Just $ SecKey fp
            else return $ Nothing
    | otherwise = Nothing

-- | Convert signature to a normalized lower-S form. Boolean value 'True'
-- indicates that the signature was normalized, 'False' indicates that it was
-- already normal.
normalizeSig :: Sig -> (Sig, Bool)
normalizeSig (Sig fg) = unsafePerformIO $ do
    fg' <- mallocForeignPtr
    ret <- withForeignPtr fg $ \pg -> withForeignPtr fg' $ \pg' ->
        ecdsa_signature_normalize ctx pg' pg
    return (Sig fg', isSuccess ret)

-- | Create internal tweak data from 32-byte 'ByteString'.
tweak :: ByteString -> Maybe Tweak
tweak bs
    | BS.length bs == 32 = unsafePerformIO $ do
        fp <- mallocForeignPtr
        withForeignPtr fp $ flip poke (Tweak32 bs)
        return $ Just $ Tweak fp
    | otherwise = Nothing

-- | Get 32-byte secret key.
getSecKey :: SecKey -> ByteString
getSecKey (SecKey fk) = getSecKey32 $ unsafePerformIO $ withForeignPtr fk peek

-- | Get 64-byte internal public key representation.
getPubKey :: PubKey -> ByteString
getPubKey (PubKey fp) = getPubKey64 $ unsafePerformIO $ withForeignPtr fp peek

-- | Get 32-byte message.
getMsg :: Msg -> ByteString
getMsg (Msg fm) = getMsg32 $ unsafePerformIO $ withForeignPtr fm $ peek

-- | Get 32-byte tweak.
getTweak :: Tweak -> ByteString
getTweak (Tweak ft) = getTweak32 $ unsafePerformIO $ withForeignPtr ft $ peek

-- | Read DER-encoded public key.
importPubKey :: ByteString -> Maybe PubKey
importPubKey bs = unsafePerformIO $ do
    useByteString bs $ \(b, l) -> do
        fp <- mallocForeignPtr
        ret <- withForeignPtr fp $ \p -> ec_pubkey_parse ctx p b l
        if isSuccess ret then return $ Just $ PubKey fp else return Nothing

-- | Encode public key as DER.  First argument 'True' for compressed output.
exportPubKey :: Bool -> PubKey -> ByteString
exportPubKey compress (PubKey pub) = unsafePerformIO $
    withForeignPtr pub $ \p -> alloca $ \l -> allocaBytes z $ \o -> do
        poke l (fromIntegral z)
        ret <- ec_pubkey_serialize ctx o l p c
        unless (isSuccess ret) $ error "could not serialize public key"
        n <- peek l
        packByteString (o, n)
  where
    c = if compress then compressed else uncompressed
    z = if compress then 33 else 65

-- | Get compact signature.
exportCompactSig :: Sig -> CompactSig
exportCompactSig (Sig fg) = unsafePerformIO $
    withForeignPtr fg $ \pg -> alloca $ \pc -> do
        ret <- ecdsa_signature_serialize_compact ctx pc pg
        unless (isSuccess ret) $ error "Could not obtain compact signature"
        peek pc

-- | Read DER-encoded signature.
importCompactSig :: CompactSig -> Maybe Sig
importCompactSig c = unsafePerformIO $ alloca $ \pc -> do
    poke pc c
    fg <- mallocForeignPtr
    ret <- withForeignPtr fg $ \pg -> ecdsa_signature_parse_compact ctx pg pc
    if isSuccess ret then return $ Just $ Sig fg else return Nothing

-- | Read DER-encoded signature.
importSig :: ByteString -> Maybe Sig
importSig bs = unsafePerformIO $
    useByteString bs $ \(b, l) -> do
        fg <- mallocForeignPtr
        ret <- withForeignPtr fg $ \g -> ecdsa_signature_parse_der ctx g b l
        if isSuccess ret then return $ Just $ Sig fg else return Nothing

-- | Relaxed DER parsing. Allows certain DER errors and violations.
laxImportSig :: ByteString -> Maybe Sig
laxImportSig bs = unsafePerformIO $
    useByteString bs $ \(b, l) -> do
        fg <- mallocForeignPtr
        ret <- withForeignPtr fg $ \g -> lax_der_parse ctx g b l
        if isSuccess ret then return $ Just $ Sig fg else return Nothing

-- | Encode signature as DER.
exportSig :: Sig -> ByteString
exportSig (Sig fg) = unsafePerformIO $
    withForeignPtr fg $ \g -> alloca $ \l -> allocaBytes 72 $ \o -> do
        poke l 72
        ret <- ecdsa_signature_serialize_der ctx o l g
        unless (isSuccess ret) $ error "could not serialize signature"
        n <- peek l
        packByteString (o, n)

-- | Verify message signature. 'True' means that the signature is correct.
verifySig :: PubKey -> Sig -> Msg -> Bool
verifySig (PubKey fp) (Sig fg) (Msg fm) = unsafePerformIO $
    withForeignPtr fp $ \p -> withForeignPtr fg $ \g ->
        withForeignPtr fm $ \m -> isSuccess <$> ecdsa_verify ctx g m p

-- | Sign message using secret key.
signMsg :: SecKey -> Msg -> Sig
signMsg (SecKey fk) (Msg fm) = unsafePerformIO $
    withForeignPtr fk $ \k -> withForeignPtr fm $ \m -> do
        fg <- mallocForeignPtr
        ret <- withForeignPtr fg $ \g -> ecdsa_sign ctx g m k nullFunPtr nullPtr
        unless (isSuccess ret) $ error "could not sign message"
        return $ Sig fg

-- | Obtain public key from secret key.
pubKey :: SecKey -> PubKey
pubKey (SecKey fk) = unsafePerformIO $
    withForeignPtr fk $ \k -> do
        fp <- mallocForeignPtr
        ret <- withForeignPtr fp $ \p -> ec_pubkey_create ctx p k
        unless (isSuccess ret) $ error "could not compute public key"
        return $ PubKey fp

-- | Read BER-encoded secret key.
importSecKey :: ByteString -> Maybe SecKey
importSecKey bs = unsafePerformIO $
    useByteString bs $ \(b, l) -> do
        fk <- mallocForeignPtr
        ret <- withForeignPtr fk $ \k -> ec_privkey_import ctx k b l
        if isSuccess ret then return $ Just $ SecKey fk else return Nothing

-- | Encode secret key as BER.  First argument 'True' for compressed output.
exportSecKey :: Bool -> SecKey -> ByteString
exportSecKey compress (SecKey fk) = unsafePerformIO $
    withForeignPtr fk $ \k -> alloca $ \l -> allocaBytes 279 $ \o -> do
        poke l 279
        ret <- ec_privkey_export ctx o l k c
        unless (isSuccess ret) $ error "could not export secret key"
        n <- peek l
        packByteString (o, n)
  where
    c = if compress then compressed else uncompressed

-- | Add tweak to secret key using ECDSA addition.
tweakAddSecKey :: SecKey -> Tweak -> Maybe SecKey
tweakAddSecKey (SecKey fk) (Tweak ft) = unsafePerformIO $
    withForeignPtr fk $ \k -> withForeignPtr ft $ \t -> do
        fk' <- mallocForeignPtr
        ret <- withForeignPtr fk' $ \k' ->  do
            key <- peek k
            poke k' key
            ec_privkey_tweak_add ctx k' t
        if isSuccess ret then return $ Just $ SecKey fk' else return Nothing

-- | Multiply secret key by tweak using ECDSA multiplication.
tweakMulSecKey :: SecKey -> Tweak -> Maybe SecKey
tweakMulSecKey (SecKey fk) (Tweak ft) = unsafePerformIO $
    withForeignPtr fk $ \k -> withForeignPtr ft $ \t -> do
        fk' <- mallocForeignPtr
        ret <- withForeignPtr fk' $ \k' ->  do
            key <- peek k
            poke k' key
            ec_privkey_tweak_mul ctx k' t
        if isSuccess ret then return $ Just $ SecKey fk' else return Nothing

-- | Perform ECDSA addition between the public key point and the point obtained
-- by multiplying the tweak scalar by the curve generator.
tweakAddPubKey :: PubKey -> Tweak -> Maybe PubKey
tweakAddPubKey (PubKey fp) (Tweak ft) = unsafePerformIO $
    withForeignPtr fp $ \p -> withForeignPtr ft $ \t -> do
        fp' <- mallocForeignPtr
        ret <- withForeignPtr fp' $ \p' ->  do
            pub <- peek p
            poke p' pub
            ec_pubkey_tweak_add ctx p' t
        if isSuccess ret then return $ Just $ PubKey fp' else return Nothing

-- | Perform ECDSA multiplication between the public key point and the point
-- obtained by multiplying the tweak scalar by the curve generator.
tweakMulPubKey :: PubKey -> Tweak -> Maybe PubKey
tweakMulPubKey (PubKey fp) (Tweak ft) = unsafePerformIO $
    withForeignPtr fp $ \p -> withForeignPtr ft $ \t -> do
        fp' <- mallocForeignPtr
        ret <- withForeignPtr fp' $ \p' ->  do
            pub <- peek p
            poke p' pub
            ec_pubkey_tweak_mul ctx p' t
        if isSuccess ret then return $ Just $ PubKey fp' else return Nothing

-- | Add multiple public keys together using ECDSA addition.
combinePubKeys :: [PubKey] -> Maybe PubKey
combinePubKeys pubs = unsafePerformIO $ pointers [] pubs $ \ps ->
    allocaArray (length ps) $ \a -> do
        pokeArray a ps
        fp <- mallocForeignPtr
        ret <- withForeignPtr fp $ \p ->
            ec_pubkey_combine ctx p a (fromIntegral $ length ps)
        if isSuccess ret
            then return $ Just $ PubKey fp
            else return Nothing
  where
    pointers ps [] f = f ps
    pointers ps (PubKey fp : pubs') f =
        withForeignPtr fp $ \p -> pointers (p:ps) pubs' f

instance Arbitrary Msg where
    arbitrary = gen_msg
      where
        valid_bs = bs_gen `suchThat` isJust
        bs_gen = (msg . BS.pack) <$> sequence (replicate 32 arbitrary)
        gen_msg = fromJust <$> valid_bs

instance Arbitrary SecKey where
    arbitrary = gen_key
      where
        valid_bs = bs_gen `suchThat` isJust
        bs_gen = (secKey . BS.pack) <$> sequence (replicate 32 arbitrary)
        gen_key = fromJust <$> valid_bs

instance Arbitrary PubKey where
    arbitrary = do
        key <- arbitrary
        return $ pubKey key