    -- |
    -- Module      : Crypto.ECC
    -- License     : BSD-style
    -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
    -- Stability   : experimental
    -- Portability : unknown
    --
    -- Elliptic Curve Cryptography
    --
    {-# LANGUAGE GeneralizedNewtypeDeriving #-}
    {-# LANGUAGE TypeFamilies #-}
    {-# LANGUAGE ScopedTypeVariables #-}
    module Crypto.ECC
        ( Curve_P256R1(..)
        , Curve_P384R1(..)
        , Curve_P521R1(..)
        , Curve_X25519(..)
        , Curve_X448(..)
        , EllipticCurve(..)
        , EllipticCurveDH(..)
        , EllipticCurveArith(..)
        , KeyPair(..)
        , SharedSecret(..)
        ) where

    import qualified Crypto.PubKey.ECC.P256 as P256
    import qualified Crypto.ECC.Simple.Types as Simple
    import qualified Crypto.ECC.Simple.Prim as Simple
    import           Crypto.Random
    import           Crypto.Error
    import           Crypto.Internal.Proxy
    import           Crypto.Internal.Imports
    import           Crypto.Internal.ByteArray (ByteArray, ByteArrayAccess, ScrubbedBytes)
    import qualified Crypto.Internal.ByteArray as B
    import           Crypto.Number.Serialize (i2ospOf_, os2ip)
    import qualified Crypto.PubKey.Curve25519 as X25519
    import qualified Crypto.PubKey.Curve448 as X448
    import           Data.Function (on)
    import           Data.ByteArray (convert)

    -- | An elliptic curve key pair composed of the private part (a scalar), and
    -- the associated point.
    data KeyPair curve = KeyPair
        { keypairGetPublic  :: !(Point curve)
        , keypairGetPrivate :: !(Scalar curve)
        }

    newtype SharedSecret = SharedSecret ScrubbedBytes
        deriving (Eq, ByteArrayAccess)

    class EllipticCurve curve where
        -- | Point on an Elliptic Curve
        type Point curve  :: *

        -- | Scalar in the Elliptic Curve domain
        type Scalar curve :: *

        -- | Generate a new random scalar on the curve.
        -- The scalar will represent a number between 1 and the order of the curve non included
        curveGenerateScalar :: MonadRandom randomly => proxy curve -> randomly (Scalar curve)

        -- | Generate a new random keypair
        curveGenerateKeyPair :: MonadRandom randomly => proxy curve -> randomly (KeyPair curve)

        -- | Get the curve size in bits
        curveSizeBits :: proxy curve -> Int

        -- | Encode a elliptic curve point into binary form
        encodePoint :: ByteArray bs => proxy curve -> Point curve -> bs

        -- | Try to decode the binary form of an elliptic curve point
        decodePoint :: ByteArray bs => proxy curve -> bs -> CryptoFailable (Point curve)

    class EllipticCurve curve => EllipticCurveDH curve where
        -- | Generate a Diffie hellman secret value.
        --
        -- This is generally just the .x coordinate of the resulting point, that
        -- is not hashed.
        --
        -- use `pointSmul` to keep the result in Point format.
        ecdh :: proxy curve -> Scalar curve -> Point curve -> SharedSecret

    class EllipticCurve curve => EllipticCurveArith curve where
        -- | Add points on a curve
        pointAdd :: proxy curve -> Point curve -> Point curve -> Point curve

        -- | Scalar Multiplication on a curve
        pointSmul :: proxy curve -> Scalar curve -> Point curve -> Point curve

    --   -- | Scalar Inverse
    --   scalarInverse :: Scalar curve -> Scalar curve

    -- | P256 Curve
    --
    -- also known as P256
    data Curve_P256R1 = Curve_P256R1

    instance EllipticCurve Curve_P256R1 where
        type Point Curve_P256R1 = P256.Point
        type Scalar Curve_P256R1 = P256.Scalar
        curveSizeBits _ = 256
        curveGenerateScalar _ = P256.scalarGenerate
        curveGenerateKeyPair _ = toKeyPair <$> P256.scalarGenerate
          where toKeyPair scalar = KeyPair (P256.toPoint scalar) scalar
        encodePoint _ p = mxy
          where
            mxy :: forall bs. ByteArray bs => bs
            mxy = B.concat [uncompressed, xy]
              where
                uncompressed, xy :: bs
                uncompressed = B.singleton 4
                xy = P256.pointToBinary p
        decodePoint _ mxy = case B.uncons mxy of
            Nothing -> CryptoFailed $ CryptoError_PointSizeInvalid
            Just (m,xy)
                -- uncompressed
                | m == 4 -> P256.pointFromBinary xy
                | otherwise -> CryptoFailed $ CryptoError_PointFormatInvalid

    instance EllipticCurveArith Curve_P256R1 where
        pointAdd  _ a b = P256.pointAdd a b
        pointSmul _ s p = P256.pointMul s p

    instance EllipticCurveDH Curve_P256R1 where
        ecdh _ s p = SharedSecret $ P256.pointDh s p

    data Curve_P384R1 = Curve_P384R1

    instance EllipticCurve Curve_P384R1 where
        type Point Curve_P384R1 = Simple.Point Simple.SEC_p384r1
        type Scalar Curve_P384R1 = Simple.Scalar Simple.SEC_p384r1
        curveSizeBits _ = 384
        curveGenerateScalar _ = Simple.scalarGenerate
        curveGenerateKeyPair _ = toKeyPair <$> Simple.scalarGenerate
          where toKeyPair scalar = KeyPair (Simple.pointBaseMul scalar) scalar
        encodePoint _ point = encodeECPoint point
        decodePoint _ bs = decodeECPoint bs

    instance EllipticCurveArith Curve_P384R1 where
        pointAdd _ a b = Simple.pointAdd a b
        pointSmul _ s p = Simple.pointMul s p

    instance EllipticCurveDH Curve_P384R1 where
        ecdh _ s p = SharedSecret $ i2ospOf_ (curveSizeBytes prx) x
          where
            prx = Proxy :: Proxy Curve_P384R1
            Simple.Point x _ = pointSmul prx s p

    data Curve_P521R1 = Curve_P521R1

    instance EllipticCurve Curve_P521R1 where
        type Point Curve_P521R1 = Simple.Point Simple.SEC_p521r1
        type Scalar Curve_P521R1 = Simple.Scalar Simple.SEC_p521r1
        curveSizeBits _ = 521
        curveGenerateScalar _ = Simple.scalarGenerate
        curveGenerateKeyPair _ = toKeyPair <$> Simple.scalarGenerate
          where toKeyPair scalar = KeyPair (Simple.pointBaseMul scalar) scalar
        encodePoint _ point = encodeECPoint point
        decodePoint _ bs = decodeECPoint bs

    instance EllipticCurveArith Curve_P521R1 where
        pointAdd _ a b = Simple.pointAdd a b
        pointSmul _ s p = Simple.pointMul s p

    instance EllipticCurveDH Curve_P521R1 where
        ecdh _ s p = SharedSecret $ i2ospOf_ (curveSizeBytes prx) x
          where
            prx = Proxy :: Proxy Curve_P521R1
            Simple.Point x _ = pointSmul prx s p

    data Curve_X25519 = Curve_X25519

    instance EllipticCurve Curve_X25519 where
        type Point Curve_X25519 = X25519.PublicKey
        type Scalar Curve_X25519 = X25519.SecretKey
        curveSizeBits _ = 255
        curveGenerateScalar _ = X25519.generateSecretKey
        curveGenerateKeyPair _ = do
            s <- X25519.generateSecretKey
            return $ KeyPair (X25519.toPublic s) s
        encodePoint _ p = B.convert p
        decodePoint _ bs = X25519.publicKey bs

    instance EllipticCurveDH Curve_X25519 where
        ecdh _ s p = SharedSecret $ convert secret
          where secret = X25519.dh p s

    data Curve_X448 = Curve_X448

    instance EllipticCurve Curve_X448 where
        type Point Curve_X448 = X448.PublicKey
        type Scalar Curve_X448 = X448.SecretKey
        curveSizeBits _ = 448
        curveGenerateScalar _ = X448.generateSecretKey
        curveGenerateKeyPair _ = do
            s <- X448.generateSecretKey
            return $ KeyPair (X448.toPublic s) s
        encodePoint _ p = B.convert p
        decodePoint _ bs = X448.publicKey bs

    instance EllipticCurveDH Curve_X448 where
        ecdh _ s p = SharedSecret $ convert secret
          where secret = X448.dh p s

    encodeECPoint :: forall curve bs . (Simple.Curve curve, ByteArray bs) => Simple.Point curve -> bs
    encodeECPoint Simple.PointO      = error "encodeECPoint: cannot serialize point at infinity"
    encodeECPoint (Simple.Point x y) = B.concat [uncompressed,xb,yb]
      where
        size = Simple.curveSizeBytes (Proxy :: Proxy curve)
        uncompressed, xb, yb :: bs
        uncompressed = B.singleton 4
        xb = i2ospOf_ size x
        yb = i2ospOf_ size y

    decodeECPoint :: (Simple.Curve curve, ByteArray bs) => bs -> CryptoFailable (Simple.Point curve)
    decodeECPoint mxy = case B.uncons mxy of
        Nothing     -> CryptoFailed $ CryptoError_PointSizeInvalid
        Just (m,xy)
            -- uncompressed
            | m == 4 ->
                let siz = B.length xy `div` 2
                    (xb,yb) = B.splitAt siz xy
                    x = os2ip xb
                    y = os2ip yb
                 in Simple.pointFromIntegers (x,y)
            | otherwise -> CryptoFailed $ CryptoError_PointFormatInvalid

    curveSizeBytes :: EllipticCurve c => Proxy c -> Int
    curveSizeBytes proxy = (curveSizeBits proxy + 7) `div` 8