{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns #-}
{-# LANGUAGE NamedFieldPuns #-}
module Crypto.Tox
    ( PublicKey
    , publicKey
    , getPublicKey
    , putPublicKey
    , SecretKey
    , generateSecretKey
    , toPublic
    , SymmetricKey(..)
    , TransportCrypto(..)
    , SecretsCache
    , newSecretsCache
    , Encrypted
    , Encrypted8(..)
    , type (∘)(..)
    , Asymm(..)
    , getAsymm
    , getAliasedAsymm
    , putAsymm
    , putAliasedAsymm
    , Plain
    , encodePlain
    , decodePlain
    -- , computeSharedSecret
    , lookupSharedSecret
    , encrypt
    , decrypt
    , Nonce8(..)
    , Nonce24(..)
    , incrementNonce24
    , addtoNonce24
    , Nonce32(..)
    , getRemainingEncrypted
    , putEncrypted
    , Auth
    , Sized(..)
    , Size(..)
    , State(..)
    , zeros32
    , zeros24
    , decryptSymmetric
    , encryptSymmetric
    , encodeSecret
    , decodeSecret
    ) where

import Control.Arrow
import Control.Monad
import qualified Crypto.Cipher.ChaChaPoly1305 as Symmetric
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 Data.Bits
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.Base64 as Base64
import qualified Data.ByteString.Char8  as C8
import Data.Data
import Data.Functor.Contravariant
#if MIN_VERSION_base(4,9,1)
import Data.Kind
#else
import GHC.Exts (Constraint)
#endif
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)
import Network.Socket (SockAddr)
import GHC.Exts (Word(..))
import GHC.Prim
import Data.Word64Map (fitsInInt)
import Data.MinMaxPSQ (MinMaxPSQ')
import qualified Data.MinMaxPSQ as MM
import Data.Time.Clock.POSIX
import Data.Hashable
import System.IO.Unsafe (unsafePerformIO)

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

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

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

infixr ∘

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 "Crypto.Tox" [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 :: Size a of
        ConstSize n -> ConstSize $ n + 16
        VarSize _   -> VarSize $ \(Encrypted bs) -> B.length bs

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 (SymmetricKey symmkey) (Nonce24 n24) (Encrypted bs) = do
    let sym_nonce_bytes = B.take 12 n24
        (mac, bs'')     = B.splitAt 16 bs
    symm <- left show . Cryptonite.eitherCryptoError $ do
        sym_nonce <- Symmetric.nonce12 sym_nonce_bytes
        Symmetric.initialize symmkey sym_nonce
    let (ds, symm') = Symmetric.decrypt bs'' symm
        auth = Symmetric.finalize symm'
    if BA.convert auth /= mac
        then Left "symmetricDecipher: Auth fail."
        else return $ Plain ds

encryptSymmetric :: SymmetricKey -> Nonce24 -> Plain s x -> Encrypted x
encryptSymmetric (SymmetricKey symmkey) (Nonce24 n24) (Plain bs) = Encrypted es
 where
    Cryptonite.CryptoPassed es = do
        sym_nonce <- Symmetric.nonce12 (BA.take 12 n24)
        symm <- Symmetric.initialize symmkey sym_nonce
        let (rpath_bs, symm') = Symmetric.encrypt bs symm
            auth = Symmetric.finalize symm' -- 16 bytes
        return (BA.convert auth `BA.append` rpath_bs)


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 ->
    let -- 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
    in State hash crypt
 where
    -- diffie helman
    shared = ecdh (Proxy :: Proxy Curve_X25519) sk recipient
    -- shared secret XSalsa key
    k = hsalsa20 shared zeros24

unsafeFirstWord64 :: ByteArrayAccess ba => ba -> Word64
unsafeFirstWord64 ba = unsafePerformIO $ BA.withByteArray ba peek

instance Hashable PublicKey where
    hashWithSalt salt pk = hashWithSalt salt (unsafeFirstWord64 pk)

instance Hashable SecretKey where
    hashWithSalt salt sk = hashWithSalt salt (unsafeFirstWord64 sk)

instance Ord PublicKey where compare = unsafeCompare32Bytes
instance Ord SecretKey where compare = unsafeCompare32Bytes

unsafeCompare32Bytes :: (ByteArrayAccess ba, ByteArrayAccess bb)
                        => ba -> bb -> Ordering
unsafeCompare32Bytes ba bb =
    unsafePerformIO $ BA.withByteArray ba
                    $ \pa -> BA.withByteArray bb
                    $ \pb -> unsafeCompare32Bytes' 3 pa pb

unsafeCompare32Bytes' :: Int -> Ptr Word64 -> Ptr Word64 -> IO Ordering
unsafeCompare32Bytes' !n !pa !pb = do
    a <- peek pa
    b <- peek pb
    case compare a b of
        EQ     -> if n == 0
                    then return EQ
                    else unsafeCompare32Bytes' (n - 1)
                                               (pa `plusPtr` 8)
                                               (pb `plusPtr` 8)
        result -> return result



lookupSharedSecret :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> IO State
lookupSharedSecret TransportCrypto{secretsCache} sk recipient nonce = do
    now <- getPOSIXTime
    atomically $ do
    mm <- readTVar $ sharedSecret secretsCache
    case MM.lookup' recipient mm of
        Nothing -> do
            let miss = computeSharedSecret sk recipient
            writeTVar (sharedSecret secretsCache)
                      (MM.insertTake' 160 recipient (MM.singleton' sk miss (Down now)) (Down now) mm)
            return $ miss nonce
        Just (stamp,smm) -> do
            let (r,v) = case MM.lookup' sk smm of
                    Nothing | let miss = computeSharedSecret sk recipient
                                      -> (miss, MM.insertTake' 3 sk miss (Down now) smm)
                    Just (stamp2,hit) -> (hit , MM.insert' sk hit (Down now) smm)
            writeTVar (sharedSecret secretsCache)
                      (MM.insertTake' 160 recipient v (Down now) mm)
            return $ r nonce


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)

addtoNonce24 :: Nonce24 -> Word -> IO Nonce24
addtoNonce24 (Nonce24 n24) n = Nonce24 <$> BA.copy n24 init
    where
        init :: Ptr Word -> IO ()
        init ptr | fitsInInt (Proxy :: Proxy Word64) = do
            let frmBE64 = fromIntegral . fromBE64 . fromIntegral
                tBE64   = fromIntegral . toBE64   . fromIntegral
                !(W# input) = n
            W# w1 <- frmBE64 <$> peek ptr
            W# w2 <- frmBE64 <$> peekElemOff ptr 1
            W# w3 <- frmBE64 <$> peekElemOff ptr 2
            let (# overflw,  sum   #) = plusWord2# w3 input
                (# overflw', sum'  #) = plusWord2# w2 overflw
                (# discard,  sum'' #) = plusWord2# w1 overflw'
            poke        ptr   $ tBE64 (W# sum'')
            pokeElemOff ptr 1 $ tBE64 (W# sum')
            pokeElemOff ptr 2 $ tBE64 (W# sum)

        init ptr | fitsInInt (Proxy :: Proxy Word32) = do
            let frmBE32 = fromIntegral . fromBE32 . fromIntegral
                tBE32   = fromIntegral . toBE32   . fromIntegral
                !(W# input) = n
            W# w1 <- frmBE32 <$> peek ptr
            W# w2 <- frmBE32 <$> peekElemOff ptr 1
            W# w3 <- frmBE32 <$> peekElemOff ptr 2
            W# w4 <- frmBE32 <$> peekElemOff ptr 3
            W# w5 <- frmBE32 <$> peekElemOff ptr 4
            W# w6 <- frmBE32 <$> peekElemOff ptr 5
            let (# overflw_,   sum_     #) = plusWord2# w6 input
                (# overflw__,  sum__    #) = plusWord2# w5 overflw_
                (# overflw___, sum___   #) = plusWord2# w6 overflw__
                (# overflw,  sum   #) = plusWord2# w3 overflw___
                (# overflw', sum'  #) = plusWord2# w2 overflw
                (# discard,  sum'' #) = plusWord2# w1 overflw'
            poke        ptr   $ tBE32 (W# sum'')
            pokeElemOff ptr 1 $ tBE32 (W# sum')
            pokeElemOff ptr 2 $ tBE32 (W# sum)
            pokeElemOff ptr 3 $ tBE32 (W# sum___)
            pokeElemOff ptr 4 $ tBE32 (W# sum__)
            pokeElemOff ptr 5 $ tBE32 (W# sum_)
        init _ = error "incrementNonce24: I only support 64 and 32 bits"

incrementNonce24 :: Nonce24 -> IO Nonce24
incrementNonce24 nonce24 = addtoNonce24 nonce24 1

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

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

bin2base64 :: ByteArrayAccess bs => bs -> String
bin2base64 = C8.unpack . Base64.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 Show Nonce32 where
    showsPrec d nonce = mappend $ bin2base64 nonce

instance Read Nonce32 where
    readsPrec _ str = either (const []) id $ do
        let (ds,ss) = Prelude.splitAt 43 str
        ss' <- case ss of
                '=':xs -> Right xs -- optional terminating '='
                _      -> Right ss
        bs <- Base64.decode (C8.pack $ ds ++ ['='])
        guard $ B.length bs == 32
        return [ (Nonce32 bs, ss') ]

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 Asymm a = Asymm
    { senderKey :: PublicKey
    , asymmNonce :: Nonce24
    , asymmData :: a
    }
 deriving (Functor,Foldable,Traversable, Show)

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

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

putAsymm :: Serialize a => Asymm a -> Put
putAsymm (Asymm key nonce dta) = putPublicKey key >> put nonce >> put dta

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

putAliasedAsymm :: Serialize a => Asymm a -> Put
putAliasedAsymm (Asymm key nonce dta) = put nonce >> putPublicKey key >> put dta

data SecretsCache = SecretsCache
    { sharedSecret :: TVar (MinMaxPSQ' PublicKey
                                       (Down POSIXTime)
                                       (MinMaxPSQ' SecretKey (Down POSIXTime) (Nonce24 -> State)))
    }

newSecretsCache :: IO SecretsCache
newSecretsCache = atomically (SecretsCache <$> newTVar MM.empty)


newtype SymmetricKey = SymmetricKey ByteString

data TransportCrypto = TransportCrypto
    { transportSecret :: SecretKey
    , transportPublic :: PublicKey
    , onionAliasSecret :: SecretKey
    , onionAliasPublic :: PublicKey
    , rendezvousSecret :: SecretKey
    , rendezvousPublic :: PublicKey
    , transportSymmetric :: STM SymmetricKey
    , transportNewNonce :: STM Nonce24
    , userKeys :: TVar [(SecretKey,PublicKey)]
    , pendingCookies :: TVar [(SockAddr, (Int, PublicKey))]
    , secretsCache :: SecretsCache
    }

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

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

-- 32 bytes -> 42 base64 digits.
--
encodeSecret :: SecretKey -> Maybe C8.ByteString
encodeSecret k = do
    (a,bs) <- BA.uncons (BA.convert k)
    -- Bytes
    -- 1   31
    -- a | bs
    (cs,c) <- unsnoc bs
    -- Bytes
    --  1   30   1
    --  a | cs | c
    --
    -- Based on the following pasted from the generateSecretKey function:
    --
    --   tweakToSecretKey :: ScrubbedBytes -> SecretKey
    --   tweakToSecretKey bin = SecretKey $ B.copyAndFreeze bin $ \inp -> do
    --         modifyByte inp 0 (\e0 -> e0 .&. 0xf8)
    --         modifyByte inp 31 (\e31 -> (e31 .&. 0x7f) .|. 0x40)
    --
    -- We know the following holds:
    --   a == a .&. 0xf8
    --   c == (c .&. 0x7f) .|. 0x40
    --
    -- Therefore, there are 5 reserved bits:
    --   a := aaaa a000
    --   c := 01dd cccc
    --
    -- That gives us 256 - 5 = 251 bits to encode.
    -- 42 * 6 = 252
    --
    let -- We'll reserve the first bit as zero so that the encoded
        -- key starts with a digit between A and f.  Other digits will be
        -- arbitrary.
        --
        -- The middle 30 bytes will be encoded as is from the source byte
        -- string (cs). It remains to compute the first (a') and last (c')
        -- bytes.
        xs = Base64.encode $ a' `BA.cons` cs `BA.snoc` c'
        -- a' := 0aaaaadd
        a' = shiftR a 1 .|. (shiftR c 4 .&. 0x03)
        -- c' := cccc0000
        c' = shiftL c 4
    return $ BA.take 42 xs

-- 42 base64 digits. First digit should be between A and f.  The rest are
-- arbitrary.
decodeSecret :: C8.ByteString -> Maybe SecretKey
decodeSecret k64 | B.length k64 < 42 = Nothing
decodeSecret k64 = do
    xs <- either (const Nothing) Just $ Base64.decode $ B.append k64 "A="
    (a',ds) <- B.uncons $ B.take 32 xs
    (cs,c') <- B.unsnoc ds
    let c = 0x40 .|. shiftR c' 4 .|. ( 0x30 .&. shiftL a' 4)
        a = 0xf8 .&. shiftL a' 1
    case secretKey $ B.cons a cs `B.snoc` c of
        CryptoPassed x -> Just x
        _              -> Nothing