{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DeriveDataTypeable #-}
module ToxCrypto
    ( PublicKey
    , publicKey
    , SecretKey
    , Encrypted
    , Plain
    , computeSharedSecret
    , encrypt
    , decrypt
    , Nonce8
    , Nonce24
    , Nonce32
    , getRemainingEncrypted
    , putEncrypted
    , Auth
    , Sized(..)
    , Size(..)
    ) 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.Kind
import Data.Ord
import Data.Serialize
import Data.Word
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import System.Endian
import qualified Data.ByteString.Internal

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

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 = 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

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

data Size a = ConstSize Int
            | VarSize (a -> Int)

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 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

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 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