-- | This module implements Scuttlebutt's Box Stream.
--
-- For more information kindly refer the to protocol guide
-- https://ssbc.github.io/scuttlebutt-protocol-guide

module Ssb.Peer.BoxStream where

import           Protolude               hiding ( Identity )

import qualified Crypto.Hash.SHA256            as SHA256
import qualified Data.ByteString               as BS
import qualified Data.ByteString.Base64        as Base64
import           Data.Either.Combinators        ( mapLeft
                                                , mapRight
                                                )
import           Control.Concurrent.STM
import qualified Data.Serialize                as Serialize
import qualified Data.Serialize.Put            as Serialize
import qualified Network.Simple.TCP            as TCP

import           Ssb.Aux
import qualified Ssb.Identity                  as SSB
import qualified Ssb.Peer.SecretHandshake      as SH

import qualified Crypto.Saltine.Class          as Nacl
import qualified Crypto.Saltine.Core.Auth      as NaclAuth
import qualified Crypto.Saltine.Core.SecretBox as Nacl
import qualified Crypto.Saltine.Core.ScalarMult
                                               as Nacl

import           Pipes
import qualified Pipes.Prelude                 as P

-- | HeaderLength is the length of a Box Stream header in bytes
headerLength :: Int
headerLength = 34

-- | MaxBodyLength is the maximum Box Stream body length in bytes
maxBodyLength :: Int
maxBodyLength = 4096

data Header = Header
    { bodyLength :: Word16
    , authTag    :: AuthTag
    } deriving (Eq,Generic,Show)

instance Serialize.Serialize Header

-- | GoodBye is the message signalling the end of the Box Stream
goodByeHeader :: Header
goodByeHeader = Header
  { bodyLength = 0
  , authTag    = AuthTag $ BS.pack $ replicate authTagLength 0
  }

newHeader :: ByteString -> ByteString -> Either Text Header
newHeader authTag body = do
  let bodyLength = BS.length body
  if bodyLength >= maxLength
    then Left "body size too big"
    else return $ Header { bodyLength = fromIntegral bodyLength
                         , authTag    = AuthTag authTag
                         }
  where maxLength = fromIntegral (maxBound :: Word16)

data Message = Message
  { header :: Header
  , body   :: ByteString
  }

encryptMessage :: Nacl.Key -> Nacl.Nonce -> ByteString -> Either Text ByteString
encryptMessage key nonce buf = do
  let (authTag, ebody) = Nacl.secretboxDetached key bodyNonce buf
  header <- newHeader authTag ebody
  let eheader = Nacl.secretbox key headerNonce $ Serialize.encode header
  return $ eheader <> ebody
 where
  headerNonce = nonce
  bodyNonce   = increment nonce

-- | A breakdown of the message alignment would be nice
-- | Problem is the body size is variable

-- | TODO: deduplicate
-- | TODO: safe take, safe tail

-- TODO: handle goodbye

-- TODO: Properly describe the function w/ Nonce update
--   The decryption / update functions should return the nonce after
--   evaluation.  This increases the difficulty for describing in the language.
--   Current work around is to model the behaviour outside the function.

decryptHeader :: Nacl.Key -> Nacl.Nonce -> ByteString -> Either Text Header
decryptHeader key nonce buf = do
  let eheader = BS.take headerLength buf
  headerBuf <- withErr (errHeader eheader)
    $ Nacl.secretboxOpen key nonce eheader
  decodeByteString headerBuf :: Either Text Header
  where errHeader h = "could not decrypt header: " <> (show h :: Text)

decryptMessage :: Nacl.Key -> Nacl.Nonce -> ByteString -> Either Text ByteString
decryptMessage key nonce buf = do
  header <- decryptHeader key nonce buf

  let rest = BS.drop headerLength buf
  ebody <- withErr "message body is smaller than messages body length"
    $ takeMay (fromIntegral $ bodyLength header) rest
  withErr (errBody ebody) $ Nacl.secretboxOpenDetached
    key
    (increment nonce)
    (extractAuthTag $ authTag header)
    ebody
  where errBody b = "could not decrypt body: " <> (show b :: Text)

goodBye :: ConnState -> ByteString
goodBye state = Nacl.secretbox (key state) (nonce state) $ encodeByteString goodByeHeader

-- | clientToServerKey is the key for client to server stream encryption.
clientToServerKey :: SH.SharedSecrets -> Either Text Nacl.Key
clientToServerKey sharedSecrets = do
  secretB <- withErr errMissingB $ SH.secretB sharedSecrets
  let layer1 = SHA256.hash $ SHA256.hash
        (  SH.network sharedSecrets
        <> Nacl.encode (SH.secretab sharedSecrets)
        <> Nacl.encode (SH.secretaB sharedSecrets)
        <> Nacl.encode (SH.secretAb sharedSecrets)
        )
  let layer2 = SHA256.hash $ layer1 <> SSB.extractPublicKey secretB
  maybeToRight errKey $ Nacl.decode layer2
 where
  errMissingB = "missing shared secret B"
  errKey      = "badly formatted sodium secret box key"

-- | serverToClientKey is the key for server to client stream encryption.
serverToClientKey :: SH.SharedSecrets -> Either Text Nacl.Key
serverToClientKey sharedSecrets = do
  secretA <- withErr errMissingA $ SH.secretA sharedSecrets
  let layer1 = SHA256.hash
        (SHA256.hash
          (  SH.network sharedSecrets
          <> Nacl.encode (SH.secretab sharedSecrets)
          <> Nacl.encode (SH.secretaB sharedSecrets)
          <> Nacl.encode (SH.secretAb sharedSecrets)
          )
        )
  let layer2 = SHA256.hash (layer1 <> SSB.extractPublicKey secretA)
  maybeToRight errKey $ Nacl.decode layer2
 where
  errMissingA = "missing shared secret A"
  errKey      = "badly formatted sodium secret box key"

clientToServerNonce :: SH.SharedSecrets -> Either Text Nacl.Nonce
clientToServerNonce sharedSecrets = do
  secretb <- withErr errMissing $ SH.secretb sharedSecrets
  key     <- withErr errBadNet $ Nacl.decode $ SH.network sharedSecrets
  let auth     = NaclAuth.auth key (SSB.extractPublicKey secretb)
  let noncebuf = BS.take 24 $ Nacl.encode auth
  withErr errMsg $ Nacl.decode noncebuf
 where
  errBadNet  = "badly formatted network id"
  errMissing = "missing shared secret a"
  errMsg     = "badly formatted sodium nonce"

serverToClientNonce :: SH.SharedSecrets -> Either Text Nacl.Nonce
serverToClientNonce sharedSecrets = do
  a   <- withErr errMissing $ SH.secreta sharedSecrets
  key <- withErr errBadNet $ Nacl.decode $ SH.network sharedSecrets
  let auth     = NaclAuth.auth key (SSB.extractPublicKey a)
  let noncebuf = BS.take 24 $ Nacl.encode auth
  withErr errMsg $ Nacl.decode noncebuf
 where
  errMissing = "missing shared secret a"
  errBadNet  = "badly formatted network id"
  errHMAC    = "badly formatted HMAC"
  errMsg     = "badly formatted sodium nonce"

-- The documentation's Client / Server key and nonce terminology is replaced
-- with local/remote fields to simplify implmenentation.
data ConnState = ConnState {
      key           :: Nacl.Key
    , nonce         :: Nacl.Nonce
    , remoteKey     :: Nacl.Key
    , remoteNonce   :: Nacl.Nonce
    , buffer        :: ByteString
    , socket        :: TCP.Socket
    }

newtype Conn = Conn ((TMVar ConnState),(TMVar ConnState))

inc :: (Word64, Word64, Word64) -> (Word64, Word64, Word64)
inc (w1, w2, w3) | w3 /= maxBound = (w1, w2, w3 + 1)
                 | w2 /= maxBound = (w1, w2 + 1, 0)
                 | w1 /= maxBound = (w1 + 1, 0, 0)
                 | otherwise      = (0, 0, 0)


-- TODO : finish me
increment :: Nacl.Nonce -> Nacl.Nonce
increment arg =
  fromMaybe undefined
    $  Nacl.decode
    $  Serialize.encode w1'
    <> Serialize.encode w2'
    <> Serialize.encode w3'
 where
  noncebuf        = Nacl.encode arg
  (b1, e3)        = BS.splitAt 16 noncebuf
  (e1, e2)        = BS.splitAt 8 b1
  w3              = fromRight undefined $ Serialize.decode e3 :: Word64
  w2              = fromRight undefined $ Serialize.decode e2 :: Word64
  w1              = fromRight undefined $ Serialize.decode e1 :: Word64
  (w1', w2', w3') = inc (w1, w2, w3)


-- | TODO: update me for handling multiple encryptions

newConnState :: TCP.Socket -> SH.SharedSecrets -> Either Text ConnState
newConnState socket sharedSecrets =
  ConnState
    <$> clientToServerKey sharedSecrets
    <*> clientToServerNonce sharedSecrets
    <*> serverToClientKey sharedSecrets
    <*> serverToClientNonce sharedSecrets
    <*> Right ""
    <*> Right socket

newServerConnState :: TCP.Socket -> SH.SharedSecrets -> Either Text ConnState
newServerConnState socket sharedSecrets =
  ConnState
    <$> serverToClientKey sharedSecrets
    <*> serverToClientNonce sharedSecrets
    <*> clientToServerKey sharedSecrets
    <*> clientToServerNonce sharedSecrets
    <*> Right ""
    <*> Right socket

-- TODO: Fix underlying network functions
--  Send never seems to fail.
send :: ConnState -> ByteString -> IO (Either Text ())
send state buf = runExceptT (TCP.send (socket state) buf)

read :: ConnState -> Int -> IO (Maybe ByteString)
read state 0     = return Nothing
read state bytes = do
  buf <- TCP.recv (socket state) bytes
  case buf of
    Nothing  -> return Nothing
    Just buf ->
      if BS.length buf == bytes
        then return $ Just buf
        else fmap (buf <>) <$> read state (bytes - BS.length buf)

-- TODO: Keep connection terminology tied to peer and local.
--    Using network terminology such as 'server' can be confusing in other when
--    functions are used in other contexts.
connectClient :: TCP.Socket -> SH.SharedSecrets -> IO (Either Text Conn)
connectClient socket sharedSecrets = do
  let state = newConnState socket sharedSecrets
  case state of
    Left  err   -> return $ Left err
    Right state -> do
      rstate <- newTMVarIO state
      wstate <- newTMVarIO state
      return . return $ Conn (rstate, wstate)

connectServer :: TCP.Socket -> SH.SharedSecrets -> IO (Either Text Conn)
connectServer socket sharedSecrets = do
  let state = newServerConnState socket sharedSecrets
  case state of
    Left  err   -> return $ Left err
    Right state -> do
      rstate <- newTMVarIO state
      wstate <- newTMVarIO state
      return . return $ Conn (rstate, wstate)


disconnect :: ConnState -> IO (Either Text ())
disconnect connState = send connState $ goodBye connState

-- TODO: Find out how to avoid these stair cases

readStream'
  :: ConnState -> Int -> IO (ConnState, Either Text (Maybe ByteString))
readStream' connState bytes = if BS.length (buffer connState) >= bytes
  then do
    let (buf, rem) = BS.splitAt bytes (buffer connState)
    let connState' = connState { buffer = rem }
    return (connState', Right (Just buf))
  else do
    buf <- withErr errNoHeader <$> read connState headerLength
    let header' = buf >>= decryptHeader key' nonce'
    if header' == Right goodByeHeader
        -- TODO: Error if not enough bytes available
      then return (connState, return Nothing)
      else do
        let bodyLength' = bodyLength <$> (buf >>= decryptHeader key' nonce')
        case bodyLength' of
          Left  err         -> return $ (connState, Left err)
          Right bodyLength' -> do
            ePayload <- (withErr errNoBody)
              <$> read connState (fromIntegral bodyLength')
            case ePayload of
              Left err -> return $ (connState, Left err)
              Right payload ->
                case
                    decryptMessage key'
                                   nonce'
                                   (fromRight undefined buf <> payload)
                  of
                    Left  err     -> return (connState, Left err)
                    Right payload -> readStream'
                      (updateNonce . appendBuffer payload $ connState)
                      bytes
 where
  errNoHeader = "could not read header"
  errNoBody   = "could not read body"
  key'        = remoteKey connState
  nonce'      = remoteNonce connState
  updateNonce connState =
    connState { remoteNonce = (increment . increment) (remoteNonce connState) }
  appendBuffer buf connState = connState { buffer = buffer connState <> buf }

readStream :: Conn -> Int -> IO (Either Text (Maybe ByteString))
readStream (Conn (mVar, _)) bytes = do
  state         <- atomically $ takeTMVar mVar
  (state', ret) <- readStream' state bytes
  atomically $ putTMVar mVar state'
  return ret

sendStream' :: ConnState -> ByteString -> IO (ConnState, Either Text ())
sendStream' connState msg = do
  let eMsg = encryptMessage key' nonce' msg
  case eMsg of
    Left  err -> return (connState, Left err)
    Right buf -> do
      ret <- send connState buf
      return (updateNonce connState, ret)
 where
  key'   = key connState
  nonce' = nonce connState
  updateNonce connState =
    connState { nonce = (increment . increment) (nonce connState) }

sendStream :: Conn -> ByteString -> IO (Either Text ())
sendStream (Conn (_, mVar)) buf = do
  state         <- atomically $ takeTMVar mVar
  (state', ret) <- sendStream' state buf
  atomically $ putTMVar mVar state'
  return ret

-- | authTagLength is the AuthTag size in bytes
authTagLength = 16

newtype AuthTag = AuthTag ByteString
  deriving (Eq,Generic,Show)

instance Serialize.Serialize AuthTag where
  get = AuthTag <$> Serialize.getByteString authTagLength
  put (AuthTag buf) = Serialize.putByteString buf

extractAuthTag :: AuthTag -> ByteString
extractAuthTag (AuthTag buf) = buf

-- | Aux functions
takeMay :: Int -> ByteString -> Maybe ByteString
takeMay l arg = if BS.length arg <= l then return $ BS.take l arg else Nothing