hask-irc/hask-irc-core/Network/IRC/Bot.hs

{-# LANGUAGE FlexibleContexts, TemplateHaskell #-}

module Network.IRC.Bot
  ( In
  , sendCommandLoop
  , readMessageLoop
  , messageProcessLoop )
where

import qualified Data.Text.Format  as TF
import qualified System.Log.Logger as HSL

import ClassyPrelude
import Control.Concurrent.Lifted  (threadDelay)
import Control.Exception.Lifted   (evaluate)
import Control.Monad.State.Strict (get, put, evalStateT)
import Data.Time                  (addUTCTime)
import System.IO                  (hIsEOF)
import System.Timeout             (timeout)
import System.Log.Logger.TH       (deriveLoggers)

import qualified Network.IRC.Configuration as CF
import Network.IRC.MessageBus
import Network.IRC.Internal.Types
import Network.IRC.Protocol
import Network.IRC.Types
import Network.IRC.Util

$(deriveLoggers "HSL" [HSL.INFO, HSL.ERROR])

data RawIn = Line !UTCTime !Text | EOS deriving (Show, Eq)
data In    = Timeout | EOD | Msg !Message deriving (Show, Eq)

formatCommand :: (Exception e) => BotConfig -> Message -> IO ([e], [Text])
formatCommand botConfig@BotConfig { .. } message =
  map (second catMaybes . partitionEithers)
  . forM (defaultCommandFormatter : cmdFormatters) $ \formatter ->
      try . evaluate $ formatter botConfig message

parseLine :: (Exception e)
          => BotConfig -> UTCTime -> Text -> Map MessageParserId [MessagePart]
          -> IO ([e], [Message], Map MessageParserId [MessagePart])
parseLine botConfig@BotConfig { .. } time line msgParts =
  map mconcat . forM parsers $ \MessageParser { .. } -> do
    let parserMsgParts    = concat . maybeToList $ lookup msgParserId msgParts
    let parserMsgPartsMap = singletonMap msgParserId parserMsgParts
    eresult <- try . evaluate $ msgParser botConfig time line parserMsgParts
    return $ case eresult of
      Left e                              -> ([e], []       , parserMsgPartsMap)
      Right ParseReject                   -> ([] , []       , parserMsgPartsMap)
      Right (ParsePartial msgParts')      -> ([] , []       , singletonMap msgParserId msgParts')
      Right (ParseDone message msgParts') -> ([] , [message], singletonMap msgParserId msgParts')
  where
    parsers = defaultParsers ++ msgParsers

sendCommandLoop :: MessageChannel Message -> Bot -> IO ()
sendCommandLoop commandChan bot@Bot { .. } = do
  msg@(Message _ _ cmd) <- receiveMessage commandChan
  (exs, lines_) <- formatCommand botConfig msg

  forM_ exs $ \(ex :: SomeException) ->
    errorM ("Error while formatting command: " ++ show cmd ++ "\nError: " ++ show ex)

  unless (null lines_) $
    handle (\(e :: SomeException) -> do
              errorM ("Error while writing to connection: " ++ show e)
              closeMessageChannel commandChan) $
      forM_ lines_ $ \line -> do
        TF.hprint botSocket "{}\r\n" $ TF.Only line
        infoM . unpack $ "> " ++ line

  commandChanClosed <- isClosedMessageChannel commandChan
  unless commandChanClosed $
    case fromMessage cmd of
      Just QuitCmd -> closeMessageChannel commandChan
      _            -> sendCommandLoop commandChan bot

readMessageLoop :: MVar BotStatus -> MessageChannel In -> Bot -> Int -> IO ()
readMessageLoop mvBotStatus inChan Bot { .. } timeoutDelay = evalStateT loop mempty
  where
    msgPartTimeout = 10

    loop = do
      msgParts  <- get
      botStatus <- readMVar mvBotStatus
      case botStatus of
        Disconnected -> io $ closeMessageChannel inChan
        _            -> do
          msgParts' <- io $ do
            mLine <- try $ timeout timeoutDelay readLine
            case mLine of
              Left (e :: SomeException)     -> do
                errorM $ "Error while reading from connection: " ++ show e
                sendMessage inChan EOD >> return msgParts
              Right Nothing                 -> sendMessage inChan Timeout >> return msgParts
              Right (Just (Line time line)) -> do
                (exs, msgs, msgParts') <- parseLine botConfig time line msgParts

                forM_ exs $ \(ex :: SomeException) ->
                  errorM ("Error while parsing line: " ++ unpack line ++ "\nError: " ++ show ex)
                forM_ msgs $ sendMessage inChan . Msg

                return msgParts'
              Right (Just EOS)              -> sendMessage inChan EOD >> return msgParts

          limit <- io $ map (addUTCTime (- msgPartTimeout)) getCurrentTime
          put $ validMsgParts limit msgParts'
          loop
      where
        validMsgParts limit =
          foldl' (\m (k, v) -> insertWith (++) k [v] m) mempty
          . concat
          . filter ((> limit) . msgPartTime . snd . headEx . sortBy (flip $ comparing (msgPartTime . snd)))
          . groupAllOn (fst &&& msgPartTarget . snd)
          . asList . concatMap (uncurry (map . (,))) . mapToList

        readLine = do
          eof <- hIsEOF botSocket
          if eof
            then return EOS
            else mask $ \unmask -> do
              line <- map initEx . unmask $ hGetLine botSocket
              infoM . unpack $ "< " ++ line
              now <- getCurrentTime
              return $ Line now line

messageProcessLoop :: MessageChannel In -> MessageChannel Message -> IRC ()
messageProcessLoop inChan messageChan = loop 0
  where
    loop !idleFor = do
      status       <- get
      Bot { .. }   <- ask
      let nick     = botNick botConfig
      let origNick = botOrigNick botConfig
      let mpass    = CF.lookup "password" (config botConfig)

      nStatus <- io . mask_ $
        if idleFor >= (oneSec * botTimeout botConfig)
          then infoM "Timeout" >> return Disconnected
          else do
            when (status == Kicked) $
              threadDelay (5 * oneSec) >> newMessage JoinCmd >>= sendMessage messageChan

            mIn <- receiveMessage inChan
            case mIn of
              Timeout                  -> do
                idleMsg <- newMessage IdleMsg
                sendMessage messageChan idleMsg
                sendWhoisMessage nick origNick
                return Idle
              EOD                      -> infoM "Connection closed" >> return Disconnected
              Msg (msg@Message { .. }) -> do
                nStatus <- handleMsg nick origNick message mpass
                sendMessage messageChan msg
                return nStatus

      put nStatus
      case nStatus of
        Idle             -> loop (idleFor + oneSec)
        Disconnected     -> return ()
        NickNotAvailable -> return ()
        NickAvailable    -> return ()
        _                -> loop 0

      where
        sendWhoisMessage nick origNick =
          when (nick /= origNick && idleFor /= 0 && idleFor `mod` (10 * oneSec) == 0) $
            (newMessage . WhoisCmd . nickToText $ origNick) >>= sendMessage messageChan

        handleMsg nick origNick message mpass
          | Just (JoinMsg user)     <- fromMessage message, userNick user == nick =
              infoM "Joined" >> return Joined
          | Just (KickMsg { .. })   <- fromMessage message, kickedNick == nick    =
              infoM "Kicked" >> return Kicked
          | Just NickInUseMsg       <- fromMessage message                        =
              infoM "Nick already in use" >> return NickNotAvailable
          | Just (ModeMsg { .. })   <- fromMessage message, modeUser == Self      = do
              whenJust mpass $ \pass -> do
                msg <- newMessage $ PrivMsgReply (User (Nick "NickServ") "") $ "IDENTIFY " ++ pass
                sendMessage messageChan msg
              newMessage JoinCmd >>= sendMessage messageChan
              return Connected
          | Just (WhoisNoSuchNick n) <- fromMessage message, n == origNick        =
              infoM "Original nick available" >> return NickAvailable
          | otherwise                                                             =
              return Connected