{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE TypeOperators #-}
module ToxCrypto
    ( PublicKey
    , publicKey
    , getPublicKey
    , putPublicKey
    , SecretKey
    , SymmetricKey(..)
    , TransportCrypto(..)
    , Encrypted
    , Encrypted8(..)
    , type (∘)(..)
    , Assym(..)
    , getAssym
    , getAliasedAssym
    , putAssym
    , putAliasedAssym
    , Plain
    , encodePlain
    , decodePlain
    , computeSharedSecret
    , encrypt
    , decrypt
    , Nonce8(..)
    , Nonce24(..)
    , Nonce32(..)
    , getRemainingEncrypted
    , putEncrypted
    , Auth
    , Sized(..)
    , Size(..)
    , State(..)
    , zeros32
    , zeros24
    , decryptSymmetric
    , encryptSymmetric
    ) where

import qualified Crypto.Cipher.Salsa    as Salsa
import qualified Crypto.Cipher.XSalsa   as XSalsa
import Crypto.ECC.Class
import qualified Crypto.Error           as Cryptonite
import qualified Crypto.MAC.Poly1305    as Poly1305
import Crypto.PubKey.Curve25519
import qualified Data.ByteArray         as BA
         ;import Data.ByteArray         as BA (ByteArrayAccess, Bytes)
import Data.ByteString                  as B
import qualified Data.ByteString.Base16 as Base16
import qualified Data.ByteString.Char8  as C8
import Data.Data
import Data.Functor.Contravariant
import Data.Kind
import Data.Ord
import Data.Serialize as S
import Data.Word
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import System.Endian
import qualified Data.ByteString.Internal
import Control.Concurrent.STM
import Crypto.Error.Types                     (CryptoFailable (..), throwCryptoError)

-- | A 16-byte mac and an arbitrary-length encrypted stream.
newtype Encrypted a = Encrypted ByteString
 deriving (Eq,Ord,Data)

newtype Encrypted8 a = E8 (Encrypted (a,Nonce8))
 deriving Serialize

newtype (f ∘ g) x = Composed { uncomposed :: f (g x) }

newtype Auth = Auth Poly1305.Auth deriving (Eq, ByteArrayAccess)
instance Ord Auth where
    compare (Auth a) (Auth b) = comparing (BA.convert :: Poly1305.Auth -> Bytes) a b
instance Data Auth where
    gfoldl k z x = z x
    -- Well, this is a little wonky... XXX
    gunfold k z c = k (z (Auth . Poly1305.Auth . (BA.convert :: ByteString -> Bytes)))
    toConstr _ = con_Auth
    dataTypeOf _ = mkDataType "ToxMessage" [con_Auth]
con_Auth :: Constr
con_Auth = mkConstr (dataTypeOf (Auth (error "con_Auth"))) "Auth" [] Prefix
instance Serialize Auth where
    get = Auth . Poly1305.Auth . BA.convert <$> getBytes 16
    put (Auth (Poly1305.Auth bs)) = putByteString $ BA.convert bs

instance Typeable a => Show (Encrypted a) where
    show (Encrypted _) = "Encrypted "++show (typeOf (undefined :: a))

encryptedAuth :: Encrypted a -> Auth
encryptedAuth (Encrypted bs)
    | Right auth <- decode (B.take 16 bs) = auth
    | otherwise                           = error "encryptedAuth: insufficient bytes"

authAndBytes :: Encrypted a -> (Auth, ByteString)
authAndBytes (Encrypted bs) = (auth,bs')
 where
    (as,bs')   = B.splitAt 16 bs
    Right auth = decode as

-- | Info about a type's serialized length. Either the length is known
-- independently of the value, or the length depends on the value.
data Size a
    = VarSize (a -> Int)
    | ConstSize !Int
    deriving Typeable

instance Contravariant Size where
  contramap f sz = case sz of
      ConstSize n -> ConstSize n
      VarSize g   -> VarSize (\x -> g (f x))

instance Monoid (Size a) where
    ConstSize x `mappend` ConstSize y = ConstSize (x + y)
    VarSize f   `mappend` ConstSize y = VarSize $ \x -> f x + y
    ConstSize x `mappend` VarSize g   = VarSize $ \y -> x + g y
    VarSize f   `mappend` VarSize g   = VarSize $ \x -> f x + g x
    mempty = ConstSize 0


class Sized a where size :: Size a

instance Sized a => Serialize (Encrypted a) where
    get = case size :: Size a of
            VarSize _   -> Encrypted <$> (remaining >>= getBytes)
            ConstSize n -> Encrypted <$> getBytes (16 + n) -- 16 extra for Poly1305 mac
    put = putEncrypted

instance Sized a => Sized (Encrypted a) where
    size = case size of
        ConstSize n -> ConstSize $ n + 16
        VarSize f   -> VarSize $ \x -> f x + 16

instance (Sized a, Sized b) => Sized (a,b) where
    size = case (size :: Size a, size :: Size b) of
        (ConstSize a , ConstSize b) -> ConstSize $ a + b
        (VarSize f   , ConstSize b) -> VarSize $ \(a, _) -> f a + b
        (ConstSize a , VarSize g)   -> VarSize $ \(_, b) -> a + g b
        (VarSize f   , VarSize g)   -> VarSize $ \(a, b) -> f a + g b

getRemainingEncrypted :: Get (Encrypted a)
getRemainingEncrypted = Encrypted <$> (remaining >>= getBytes)

putEncrypted :: Encrypted a -> Put
putEncrypted (Encrypted bs) = putByteString bs

newtype Plain (s:: * -> Constraint) a = Plain ByteString


decodePlain :: Serialize a => Plain Serialize a -> Either String a
decodePlain (Plain bs) = decode bs

encodePlain :: Serialize a => a -> Plain Serialize a
encodePlain a = Plain $ encode a

storePlain :: Storable a => a -> IO (Plain Storable a)
storePlain a = Plain <$> BA.create (sizeOf a) (`poke` a)

retrievePlain :: Storable a => Plain Storable a -> IO a
retrievePlain (Plain bs) = BA.withByteArray bs peek

decryptSymmetric :: SymmetricKey -> Nonce24 -> Encrypted a -> Either String (Plain s a)
decryptSymmetric sym nonce e = _todo

encryptSymmetric :: SymmetricKey -> Nonce24 -> Plain s x -> Encrypted x
encryptSymmetric = _todo

data State = State Poly1305.State XSalsa.State

decrypt :: State -> Encrypted a -> Either String (Plain s a)
decrypt (State hash crypt) ciphertext
    | (a == mac) = Right (Plain m)
    | otherwise  = Left "decipherAndAuth: auth fail"
  where
    (mac, c) = authAndBytes ciphertext
    m = fst . XSalsa.combine crypt $ c
    a = Auth . Poly1305.finalize . Poly1305.update hash $ c

-- Encrypt-then-Mac:  Encrypt the cleartext, then compute the MAC on the
-- ciphertext, and prepend it to the ciphertext
encrypt :: State -> Plain s a -> Encrypted a
encrypt (State hash crypt) (Plain m) = Encrypted $ B.append (encode a) c
  where
    c = fst . XSalsa.combine crypt $ m
    a = Auth . Poly1305.finalize . Poly1305.update hash $ c

-- (Poly1305.State, XSalsa.State)
computeSharedSecret :: SecretKey -> PublicKey -> Nonce24 -> State
computeSharedSecret sk recipient nonce = State hash crypt
 where
    -- diffie helman
    shared = ecdh (Proxy :: Proxy Curve_X25519) sk recipient
    -- shared secret XSalsa key
    k = hsalsa20 shared zeros24
    -- cipher state
    st0 = XSalsa.initialize 20 k nonce
    -- Poly1305 key
    (rs, crypt) = XSalsa.combine st0 zs where Nonce32 zs = zeros32
    -- Since rs is 32 bytes, this pattern should never fail...
    Cryptonite.CryptoPassed hash = Poly1305.initialize rs

hsalsa20 :: (ByteArrayAccess t, ByteArrayAccess t1) => t1 -> t -> BA.ScrubbedBytes
hsalsa20 k n = BA.append a b
 where
    Salsa.State st = XSalsa.initialize 20 k n
    (_, as) = BA.splitAt 4 st
    (a, xs) = BA.splitAt 16 as
    (_, bs) = BA.splitAt 24 xs
    (b, _ ) = BA.splitAt 16 bs


newtype Nonce24 = Nonce24 ByteString
 deriving (Eq, Ord, ByteArrayAccess,Data)

quoted :: ShowS -> ShowS
quoted shows s = '"':shows ('"':s)

bin2hex :: ByteArrayAccess bs => bs -> String
bin2hex = C8.unpack . Base16.encode . BA.convert

instance Show Nonce24 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)

instance Sized Nonce24 where size = ConstSize 24

instance Serialize Nonce24 where
    get = Nonce24 <$> getBytes 24
    put (Nonce24 bs) = putByteString bs

newtype Nonce8 = Nonce8 Word64
 deriving (Eq, Ord, Data, Serialize)

-- Note: Big-endian to match Serialize instance.
instance Storable Nonce8 where
    sizeOf _ = 8
    alignment _ = alignment (undefined::Word64)
    peek ptr = Nonce8 . fromBE64 <$> peek (castPtr ptr)
    poke ptr (Nonce8 w) = poke (castPtr ptr) (toBE64 w)

instance Sized Nonce8 where size = ConstSize 8

instance ByteArrayAccess Nonce8 where
    length _ = 8
    withByteArray (Nonce8 w64) kont =
        allocaBytes 8 $ \p -> do
            poke (castPtr p :: Ptr Word64) $ toBE64 w64
            kont p

instance Show Nonce8 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)


newtype Nonce32 = Nonce32 ByteString
 deriving (Eq, Ord, ByteArrayAccess, Data)

instance Serialize Nonce32 where
    get = Nonce32 <$> getBytes 32
    put (Nonce32 bs) = putByteString bs

instance Sized Nonce32 where size = ConstSize 32


zeros32 :: Nonce32
zeros32 = Nonce32 $ BA.replicate 32 0

zeros24 :: ByteString
zeros24 = BA.take 24 zs where Nonce32 zs = zeros32

-- | `32`   | sender's DHT public key   |
-- | `24`   | nonce                     |
-- | `?`    | encrypted message         |
data Assym a = Assym
    { senderKey :: PublicKey
    , assymNonce :: Nonce24
    , assymData :: a
    }
 deriving (Functor,Foldable,Traversable, Show)

instance Sized a => Sized (Assym a) where
    size = case size of
        ConstSize a -> ConstSize $ a + 24 + 32
        VarSize f   -> VarSize $ \Assym { assymData = x } -> f x + 24 + 32

-- | Field order: senderKey, then nonce This is the format used by
-- Ping/Pong/GetNodes/SendNodes.
--
-- See 'getAliasedAssym' if the nonce precedes the key.
getAssym :: Serialize a => Get (Assym a)
getAssym = Assym <$> getPublicKey <*> get <*> get

putAssym :: Serialize a => Assym a -> Put
putAssym (Assym key nonce dta) = putPublicKey key >> put nonce >> put dta

-- | Field order: nonce, and then senderKey.
getAliasedAssym :: Serialize a => Get (Assym a)
getAliasedAssym = flip Assym <$> get <*> getPublicKey <*> get

putAliasedAssym :: Serialize a => Assym a -> Put
putAliasedAssym (Assym key nonce dta) = put nonce >> putPublicKey key >> put dta

newtype SymmetricKey = SymmetricKey ByteString

data TransportCrypto = TransportCrypto
    { transportSecret :: SecretKey
    , transportPublic :: PublicKey
    , transportSymmetric :: STM SymmetricKey
    , transportNewNonce :: STM Nonce24
    }

getPublicKey :: S.Get PublicKey
getPublicKey = throwCryptoError . publicKey <$> S.getBytes 32

putPublicKey :: PublicKey -> S.Put
putPublicKey bs = S.putByteString $ BA.convert bs