path: root/src/Network/Tox/Crypto/Handlers.hs

{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TupleSections #-}
module Network.Tox.Crypto.Handlers where

import Network.Tox.NodeId
import Network.Tox.Crypto.Transport
import Network.Tox.DHT.Transport (Cookie(..),CookieData(..), CookieRequest(..))
import Network.Tox.DHT.Handlers (Client, cookieRequest, cookieRequestH )
import Crypto.Tox
import Control.Concurrent.STM
import Network.Address
import qualified Data.Map.Strict as Map
import Crypto.Hash
import Control.Applicative
import Control.Monad
import Data.Time.Clock.POSIX
import qualified Data.ByteString  as B
import Data.ByteString (ByteString)
import Control.Lens
import Data.Function
import Data.Serialize             as S
import Data.Word
import GHC.Conc (unsafeIOToSTM)
import qualified Data.Set         as Set
import qualified Data.Array.Unboxed as A
import SensibleDir
import System.FilePath
import System.IO.Temp
import System.Environment
import System.Directory

-- util, todo: move to another module
maybeToEither :: Maybe b -> Either String b
maybeToEither (Just x) = Right x
maybeToEither Nothing  = Left "maybeToEither"

data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
    deriving (Eq,Ord,Show,Enum)


type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
type NetCryptoHook = IOHook NetCryptoSession CryptoData
type MsgTypeArray = A.UArray Word8 Word16


-- | Information, that may be made visible in multiple sessions, as well
-- as displayed in some way to the user via mutiple views.
--
data SessionView = SessionView
        { svNick :: TVar ByteString
        , svStatus :: TVar UserStatus
        , svStatusMsg :: TVar ByteString
        , svGroups :: TVar (Map.Map GroupChatId (Set.Set SockAddr))
        , svCacheDir :: FilePath -- ^ directory path used if the session has to use the disk for cache
                                 --   clean up only if space is needed
        , svTmpDir :: FilePath -- ^ Once off storage goes here, should clean up quickly
        , svConfigDir :: FilePath -- ^ profile related storage, etc, never clean up
        , svDownloadDir :: TVar FilePath -- ^ where to put files the user downloads
        }


data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
                                , ncTheirBaseNonce :: TVar Nonce24 -- base nonce + packet number
                                , ncMyPacketNonce :: TVar Nonce24 -- base nonce + packet number
                                , ncHandShake     :: TVar (Maybe (Handshake Encrypted))
                                , ncCookie        :: TVar (Maybe Cookie)
                                , ncTheirDHTKey   :: PublicKey
                                , ncTheirSessionPublic :: Maybe PublicKey
                                , ncSessionSecret :: SecretKey
                                , ncSockAddr      :: SockAddr
                                , ncHooks         :: TVar (Map.Map MessageType [NetCryptoHook])
                                , ncUnrecognizedHook :: TVar (MessageType -> NetCryptoHook)
                                , ncMessageTypes :: TVar MsgTypeArray -- ^ supported messages, 0 for unsupported,
                                                                      --   otherwise the messageType, some message types
                                                                      --   may not be in ncHooks yet, but they should appear
                                                                      --   here if ncUnrecognizedHook will add them to ncHooks
                                                                      --   on an as-need basis. On 0 entries, this Transport
                                                                      --   will return id in case they are handled by another Transport.
                                , ncAllSessions :: NetCryptoSessions -- ^ may be needed if one net-crypto session
                                                                     --   needs to possibly start another, as is
                                                                     --   the case in group chats
                                , ncView :: TVar SessionView
                                }

data NetCryptoSessions = NCSessions { netCryptoSessions :: TVar (Map.Map SockAddr NetCryptoSession)
                                    , transportCrypto :: TransportCrypto
                                    , defaultHooks :: Map.Map MessageType [NetCryptoHook]
                                    , defaultUnrecognizedHook :: MessageType -> NetCryptoHook
                                    , sessionView :: SessionView
                                    , msgTypeArray :: MsgTypeArray
                                    }

newSessionsState :: TransportCrypto
                 -> (MessageType -> NetCryptoHook) -- ^ default hook
                 -> Map.Map MessageType [NetCryptoHook] -- ^ all hooks, can be empty to start
                 -> IO NetCryptoSessions
newSessionsState crypto unrechook hooks = do
    x <- atomically $ newTVar Map.empty
    nick <- atomically $ newTVar B.empty
    status <- atomically $ newTVar Online
    statusmsg <- atomically $ newTVar B.empty
    grps <- atomically $ newTVar Map.empty
    pname <- getProgName
    cachedir <- sensibleCacheDirCreateIfMissing pname
    tmpdir <- (</> pname) <$> getCanonicalTemporaryDirectory
    configdir <- sensibleVarLib pname
    homedir <- getHomeDirectory
    svDownloadDir0 <- atomically $ newTVar (homedir </> "Downloads")
    return NCSessions { netCryptoSessions = x
                      , transportCrypto = crypto
                      , defaultHooks = hooks
                      , defaultUnrecognizedHook = unrechook
                      , sessionView = SessionView { svNick = nick
                                                  , svStatus = status
                                                  , svStatusMsg = statusmsg
                                                  , svGroups = grps
                                                  , svCacheDir = cachedir
                                                  , svTmpDir = tmpdir
                                                  , svConfigDir = configdir
                                                  , svDownloadDir = svDownloadDir0
                                                  }
                      , msgTypeArray = allMsgTypes -- todo make this a parameter
                      }

data HandshakeParams
        = HParam
            { hpTheirBaseNonce :: Maybe Nonce24 -- ignore and generate your own
            , hpOtherCookie :: Cookie
            , hpTheirSessionKeyPublic :: PublicKey
            , hpMySecretKey :: SecretKey
            , hpCookieRemotePubkey :: PublicKey
            , hpCookieRemoteDhtkey :: PublicKey
            }
newHandShakeData :: TransportCrypto -> Nonce24 -> HandshakeParams -> SockAddr -> IO (Maybe HandshakeData)
newHandShakeData crypto basenonce (HParam {hpOtherCookie,hpMySecretKey,hpCookieRemotePubkey,hpCookieRemoteDhtkey}) addr
    = do
    freshCookie
      <- case  nodeInfo (key2id hpCookieRemoteDhtkey) addr  of
            Right nodeinfo -> Just <$> cookieRequestH crypto nodeinfo (CookieRequest hpCookieRemotePubkey)
            Left er        -> return Nothing
    let hinit = hashInit
        Cookie n24 encrypted = hpOtherCookie
        hctx = hashUpdate hinit n24
        hctx' = hashUpdate hctx encrypted
        digest = hashFinalize hctx'
-- parameters addr {--> SockAddr -}
--    mbcookie <- case hpOtherCookie of
--                    Nothing -> case (nodeInfo hpCookieRemoteDhtkey addr) of
--                                Right nodeinfo -> cookieRequest crypto netCryptoDHTClient (toPublic hpMySecretKey) nodeinfo
--                                Left er -> return Nothing
--                    Just c -> return (Just c)

    return $
      fmap (\freshCookie' ->
            HandshakeData
                { baseNonce = basenonce
                , sessionKey = toPublic hpMySecretKey
                , cookieHash = digest
                , otherCookie = freshCookie'
                }) freshCookie

-- | called when we recieve a crypto handshake with valid cookie
freshCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> IO ()
freshCryptoSession  sessions
                    addr
                    hp@(HParam
                              { hpTheirBaseNonce = Just theirBaseNonce
                              , hpOtherCookie    = otherCookie
                              , hpTheirSessionKeyPublic = theirSessionKey
                              , hpMySecretKey        = key
                              , hpCookieRemotePubkey = remotePublicKey
                              , hpCookieRemoteDhtkey = remoteDhtPublicKey
                              }) = do
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
    ncState0 <- atomically $ newTVar Accepted
    ncTheirBaseNonce0 <- atomically $ newTVar theirBaseNonce
    n24 <- atomically $ transportNewNonce crypto
    state <- lookupSharedSecret crypto key remoteDhtPublicKey n24
    newBaseNonce <- atomically $ transportNewNonce crypto
    mbMyhandshakeData <- newHandShakeData crypto newBaseNonce hp addr
    let encodeHandshake myhandshakeData = let plain = encodePlain myhandshakeData
    --                                        state = computeSharedSecret key remoteDhtPublicKey n24
                                              encrypted = encrypt state plain
                                              in Handshake { handshakeCookie = otherCookie
                                                           , handshakeNonce = n24
                                                           , handshakeData = encrypted
                                                           }
    let myhandshake= encodeHandshake <$> mbMyhandshakeData
    ncMyPacketNonce0 <- atomically $ newTVar newBaseNonce
    ncHandShake0 <- atomically $ newTVar myhandshake
    cookie0 <- atomically $ newTVar (Just otherCookie)
    newsession <- generateSecretKey
    ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
    ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
    ncMessageTypes0 <- atomically $ newTVar (msgTypeArray sessions)
    ncView0 <- atomically $ newTVar (sessionView sessions)
    let netCryptoSession =
            NCrypto { ncState         = ncState0
                    , ncTheirBaseNonce= ncTheirBaseNonce0
                    , ncMyPacketNonce   = ncMyPacketNonce0
                    , ncHandShake     = ncHandShake0
                    , ncCookie        = cookie0
                    , ncTheirDHTKey   = remoteDhtPublicKey
                    , ncTheirSessionPublic = Just theirSessionKey
                    , ncSessionSecret = newsession
                    , ncSockAddr      = addr
                    , ncHooks         = ncHooks0
                    , ncUnrecognizedHook = ncUnrecognizedHook0
                    , ncAllSessions = sessions
                    , ncMessageTypes = ncMessageTypes0
                    , ncView = ncView0
                    }
    atomically $ modifyTVar allsessions (Map.insert addr netCryptoSession)

-- | Called when we get a handshake, but there's already a session entry.
updateCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> NetCryptoSession -> IO ()
updateCryptoSession sessions addr hp session = do
    ncState0 <- atomically $ readTVar (ncState session)
    ncTheirBaseNonce0 <- atomically $ readTVar (ncTheirBaseNonce session)
    if (ncState0 >= Accepted)
        -- If the nonce in the handshake and the dht key are both the same as
        -- the ones we have saved, assume we already handled this and this is a
        -- duplicate handshake packet, otherwise disregard everything, and
        -- refresh all state.
        --
     then when (  Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp
               || ncTheirDHTKey session /= hpCookieRemoteDhtkey hp
               ) $ freshCryptoSession sessions addr hp
     else if ( Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp)
            then freshCryptoSession sessions addr hp -- basenonce mismatch, trigger refresh
            else atomically $ writeTVar (ncState session) Accepted


cryptoNetHandler :: NetCryptoSessions -> SockAddr -> NetCrypto -> IO (Maybe (NetCrypto -> NetCrypto))
cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) nonce24 encrypted)) = do
    -- Handle Handshake Message
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
        anyRight [] f     = return $ Left "missing key"
        anyRight (x:xs) f = f x >>= either (const $ anyRight xs f) (return . Right)
    seckeys <- map fst <$> atomically (readTVar (userKeys crypto))
    symkey <- atomically $ transportSymmetric crypto
    now <- getPOSIXTime
    lr <- fmap join . sequence $ do -- Either Monad
            (CookieData cookieTime remotePubkey remoteDhtkey) <- (decodePlain =<< decryptSymmetric symkey n24 ecookie)
            Right $ do -- IO Monad
            decrypted <- anyRight seckeys $ \key -> do
                            secret <- lookupSharedSecret crypto key remotePubkey nonce24
                            return $ (key,) <$> (decodePlain =<< decrypt secret encrypted)
            return $ do -- Either Monad
            (key,HandshakeData { baseNonce, sessionKey, cookieHash, otherCookie }) <- decrypted
            -- check cookie time < 15 seconds ago
            guard (now - fromIntegral cookieTime < 15)
            -- cookie hash is valid? sha512 of ecookie
            let hinit = hashInit
                hctx = hashUpdate hinit n24
                hctx' = hashUpdate hctx ecookie
                digest = hashFinalize hctx'
            guard (cookieHash == digest)
            -- known friend?
            -- todo
            return
             HParam
                { hpTheirBaseNonce = Just baseNonce
                , hpOtherCookie    = otherCookie
                , hpTheirSessionKeyPublic =  sessionKey
                , hpMySecretKey        = key
                , hpCookieRemotePubkey = remotePubkey
                , hpCookieRemoteDhtkey = remoteDhtkey
                }
    case lr of
        Left _ -> return ()
        Right hp@(HParam
                      { hpTheirBaseNonce = Just theirBaseNonce
                      , hpOtherCookie    = otherCookie
                      , hpTheirSessionKeyPublic = theirSessionKey
                      , hpMySecretKey        = key
                      , hpCookieRemotePubkey = remotePublicKey
                      , hpCookieRemoteDhtkey = remoteDhtPublicKey
                      }) -> do
            sessionsmap <- atomically $ readTVar allsessions
            -- Do a lookup, so we can handle the update case differently
            case Map.lookup addr sessionsmap of
                Nothing -> freshCryptoSession sessions addr hp -- create new session
                Just session -> updateCryptoSession sessions addr hp session -- update existing session
    return Nothing


cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
    sessionsmap <- atomically $ readTVar allsessions
    -- Handle Encrypted Message
    case Map.lookup addr sessionsmap of
        Nothing -> return Nothing -- drop packet, we have no session
        Just session@(NCrypto {ncState, ncHooks,ncSessionSecret,ncTheirSessionPublic,ncTheirBaseNonce}) -> do
            theirBaseNonce <- atomically $ readTVar ncTheirBaseNonce
            -- Try to decrypt message
            let diff :: Word16
                diff = nonce16 - fromIntegral (last2Bytes theirBaseNonce) -- truncating to Word16
            tempNonce <- addtoNonce24 theirBaseNonce (fromIntegral diff) -- expanding to Word
            lr <- fmap join $ sequence $ do -- Either Monad --
                        pubkey <- maybeToEither ncTheirSessionPublic
                        Right $ do -- IO Monad
                        secret <- lookupSharedSecret crypto ncSessionSecret pubkey tempNonce
                        return $ decodePlain =<< decrypt secret encrypted
            case lr of
                Left _ -> return Nothing -- decryption failed, ignore packet
                Right cd@(CryptoData {bufferStart, bufferEnd, bufferData=cm}) -> do -- decryption succeeded,
                    -- TODO: Why do I need bufferStart & bufferEnd?
                    --
                    --   buffer_start = highest packet number handled + 1
                    --                , recvbuffers buffer_start
                    --
                    --   bufferEnd = sendbuffer buffer_end if lossy, otherwise packet number
                    -- update ncTheirBaseNonce if necessary
                    when (diff > 2 * dATA_NUM_THRESHOLD)$
                        atomically $ do
                            y <- readTVar ncTheirBaseNonce
                            -- all because Storable forces IO...
                            x <- unsafeIOToSTM $ addtoNonce24 y (fromIntegral dATA_NUM_THRESHOLD)
                            writeTVar ncTheirBaseNonce y
                    -- then set session confirmed,
                    atomically $ writeTVar ncState Confirmed
                    runCryptoHook session cd
    where
        last2Bytes :: Nonce24 -> Word
        last2Bytes (Nonce24 bs) = case S.decode (B.drop 22 bs) of
                                    Right n -> n
                                    _ -> error "unreachable-last2Bytes"
        dATA_NUM_THRESHOLD = 21845 -- = 65535 / 3

runCryptoHook :: NetCryptoSession -> CryptoData -> IO (Maybe (x -> x))
runCryptoHook session@(NCrypto {ncState, ncHooks,ncSessionSecret,ncTheirSessionPublic,ncTheirBaseNonce,ncMessageTypes})
              cd@(CryptoData {bufferStart, bufferEnd, bufferData=cm}) = do
    hookmap <- atomically $ readTVar ncHooks
    -- run hook
    flip fix (cd,hookmap) $ \lookupAgain (cd,hookmap) -> do
        msgTypes <- atomically $ readTVar ncMessageTypes
        let msgTyp = cd ^. messageType
            msgTypMapped16 = msgTypes A.! fromEnum8 (msgID cm)
            msgTypMapped = fromWord16 $ msgTypMapped16
        if msgTypMapped16 == 0
         then return $ Just id
         else
           case Map.lookup msgTypMapped hookmap of
                Nothing -> do -- no recognizing hook, run ncUnrecognizedHook0, loopAgain on result
                             unrecognize <- atomically $ readTVar (ncUnrecognizedHook session)
                             mbConsume <- unrecognize msgTypMapped session cd
                             case mbConsume of
                                Just f -> do
                                           -- ncUnrecognizedHook0 may have updated the hookmap
                                           hookmap' <- atomically $ readTVar ncHooks
                                           lookupAgain (f cd,hookmap')
                                Nothing -> return Nothing
                Just hooks -> flip fix (hooks,cd,msgTypMapped) $ \loop (hooks,cd,typ) -> do
                                let _ = cd :: CryptoData
                                case (hooks,cd) of
                                    ([],_) ->  return Nothing
                                    (hook:more,cd) -> do
                                        r <- hook session cd :: IO (Maybe (CryptoData -> CryptoData))
                                        case r of
                                            Just f -> let newcd = f cd
                                                          newtyp = newcd ^. messageType
                                                          in if newtyp == typ then loop (more,newcd,newtyp)
                                                                              else lookupAgain (newcd,hookmap)
                                            Nothing -> return Nothing -- message consumed

-- | construct a 'MsgTypeArray' for specified types, using their known common positions
--   in the MessageId space if they have such a thing.
mkMsgTypes :: [MessageType] -> MsgTypeArray
mkMsgTypes msgs = let zeros = A.listArray (0,255) (replicate 256 0)
                      in zeros A.// map (\x -> (toIndex x,toWord16 x)) msgs
    where
        toIndex (Msg mid) = fromIntegral . fromEnum  $ mid
        toIndex (GrpMsg KnownLossless nam) = 0x63 -- fromEnum MESSAGE_GROUPCHAT
        toIndex (GrpMsg KnownLossy    nam) = 0xC7 -- fromEnum LOSSY_GROUPCHAT

-- | Handle all Tox messages that this code base is aware of.
allMsgTypes :: MsgTypeArray
allMsgTypes = A.listArray (minBound,maxBound) (0:knownMsgs)
  where
    knownMsgs :: [Word16]
    knownMsgs =
      concat [ map (fromIntegral . fromEnum) [ PacketRequest .. KillPacket ]
             , map (const 0)                 [ 3 .. 15 ]  -- UnspecifiedPacket
             , map (const 0)                 [ 16 .. 23 ] -- MessengerLoseless
             , map (fromIntegral . fromEnum) [ ONLINE .. OFFLINE ]
             , map (const 0)                 [ 26 .. 47 ] -- MessengerLoseless
             , map (fromIntegral . fromEnum) [ NICKNAME .. TYPING ]
             , map (const 0)                 [ 52 .. 63 ] -- MessengerLoseless
             , map (fromIntegral . fromEnum) [ MESSAGE .. ACTION ]
             , map (const 0)                 [ 66 .. 68 ] -- MessengerLoseless
             , map (fromIntegral . fromEnum) [ MSI ]
             , map (const 0)                 [ 70 .. 79 ] -- MessengerLoseless
             , map (fromIntegral . fromEnum) [ FILE_SENDREQUEST .. FILE_DATA ]
             , map (const 0)                 [ 83 .. 95 ] -- MessengerLoseless
             , map (fromIntegral . fromEnum) [ INVITE_GROUPCHAT .. MESSAGE_GROUPCHAT ]
             , map (const 0)                 [ 100 .. 191 ] -- MessengerLoseless
             , map (const 0)                 [ 192 .. 198 ] -- MessengerLossy
             , map (fromIntegral . fromEnum) [ LOSSY_GROUPCHAT ]
             , map (const 0)                 [ 200 .. 255 ] -- All lossy, exept the last
             ]

-- | handles nothing
defaultCryptoDataHooks :: Map.Map MessageType [NetCryptoHook]
defaultCryptoDataHooks = Map.empty

-- | discards all unrecognized packets
defaultUnRecHook :: MessageType -> NetCryptoHook
defaultUnRecHook _ _ _ = return Nothing

-- | use to add a single hook to a specific session.
addCryptoDataHook1 :: Map.Map MessageType [NetCryptoHook] -> MessageType -> NetCryptoHook -> Map.Map MessageType [NetCryptoHook]
addCryptoDataHook1 mp typ hook = case Map.lookup typ mp of
                                    Nothing -> Map.insert typ [hook] mp
                                    Just hooks -> Map.insert typ (hook:hooks) mp