{-# LANGUAGE FlexibleInstances, DeriveDataTypeable #-}

{-| This module receives messages from external network components,
    such as, authentication portals, intrusion detection systems, vulnerability
    scanners, etc. These components are not necessarily written in Haskell or
    depend on Nettle; thus, the module acts an interface to the outside world.

    This module interacts with external components via a threaded TCP server.
    Users of this module are given a channel on which they can receive messages
    from external components.

    This module also assists in writing messenger clients that send messages to
    the server. Note that these clients need not be written in Haskell; it would
    be trivial to write a messenger client in another language. We provide a
    messenger client here for your reference and for convenience of integrating
    external Haskell components with Nettle.

    Messages are plain text; it is your responsibility to impose semantics for
    these strings and ensure that messages from different components don't
    clash. (For example, I recommend that all messages from a vulnerability
    scanner be prefixed with \"Scanner\", messages from a Web portal be prefixed
    with \"Portal\", etc.)

 -}

module Nettle.Servers.MessengerServer
    (
      Message
    , messengerServer
    , messengerClient
    ) where

import Data.Binary
import Data.Binary.Put
import Data.Binary.Get
import qualified Data.ByteString.Lazy.Char8 as B
import Control.Concurrent
import Nettle.Servers.TwoWayChannel
import Control.Exception (finally, handle, throwIO, Exception, handle)
import Data.Typeable
import Network
import Network.Socket
import Control.Monad (forever, when, unless, liftM, replicateM)
import System.IO
import System.IO.Error (IOError, isEOFError, mkIOError, eofErrorType)

type Message = String

-- Marshal a message to wire format and send it on a handle.
sendMessage hdl message =
    let writer = do putWord16be $ fromInteger $ toInteger $ length message
                    mapM_ put message
    in B.hPut hdl $ runPut writer

-- Receive a single message from a handle, or raise an IOError
receiveMessage hdl =
    do isEOF <- hIsEOF hdl
       when isEOF $ throwIO $ mkIOError eofErrorType "End of stream" (Just hdl) Nothing
       headerBytes <- B.hGet hdl 2
       let bodyLen = fromInteger $ toInteger (decode headerBytes :: Word16)
       replicateM bodyLen $ hGetChar hdl

-- Receive messages from external components and put them onto the channel
handleMessage addr hdl chan =
    do message <- receiveMessage hdl
       writeChan chan (addr, message)

trivialHandler :: IOError -> IO ()
trivialHandler e = print e >> return ()

-- Handle messages from a single client component.
handleClient hdl addr chan =
    handle trivialHandler $
           finally (sequence_ $ repeat $ handleMessage addr hdl chan)
                   (hClose hdl)

-- Set up a client connection
acceptConnection sock chan =
    do (sock', addr) <- Network.Socket.accept sock
       hdl <- socketToHandle sock' ReadWriteMode
       hSetBuffering hdl NoBuffering
       forkIO $ handleClient hdl addr chan

-- | Run a server that listens for connections from external messenger
-- components. The server is run in a separate thread; the current thread
-- returns a channel to the user where he can receive messages from the external
-- components. Messages are annotated with the client's address information.
messengerServer :: PortID -> IO (Chan (SockAddr, Message))
messengerServer port =
    do sock <- listenOn port
       chan <- newChan
       forkIO $ finally (forever $ acceptConnection sock chan) (sClose sock)
       return chan

data DeadException = DeadException
    deriving (Typeable, Show)
instance Exception DeadException

-- | Connect to a messenger server (i.e., a Nettle controller that runs
-- messengerServer) and send messages. Note that although this function
-- returns a bidirectional channel for both sending and receiving messages
-- from the server, the server currently does not send any messages.
messengerClient :: HostName -> PortID -> (Chan2 (Maybe Message) (Maybe Message) () -> IO ()) -> IO ()
messengerClient host port client =
    do
       hdl <- connectTo host port
       chan <- newChan2
       forkIO $ client chan >> kill chan ()
       forkIO $ whenDead chan >> writeChan2 chan Nothing
       handleMessage hdl $ theOtherEnd2 chan
       hClose hdl
    where handleMessage hdl ch = do message <- readChan2 ch
                                    case message of
                                         Just msg -> do sendMessage hdl msg
                                                        handleMessage hdl ch
                                         Nothing -> return ()