A simple IRC bot written in Haskell
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

Bot.hs 7.7KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186
  1. {-# LANGUAGE FlexibleContexts, TemplateHaskell #-}
  2. module Network.IRC.Bot
  3. ( In
  4. , sendCommandLoop
  5. , readMessageLoop
  6. , messageProcessLoop )
  7. where
  8. import qualified Data.Text.Format as TF
  9. import qualified System.Log.Logger as HSL
  10. import ClassyPrelude
  11. import Control.Concurrent.Lifted (threadDelay)
  12. import Control.Exception.Lifted (evaluate)
  13. import Control.Monad.State.Strict (get, put, evalStateT)
  14. import Data.Time (addUTCTime)
  15. import System.IO (hIsEOF)
  16. import System.Timeout (timeout)
  17. import System.Log.Logger.TH (deriveLoggers)
  18. import qualified Network.IRC.Configuration as CF
  19. import Network.IRC.MessageBus
  20. import Network.IRC.Internal.Types
  21. import Network.IRC.Protocol
  22. import Network.IRC.Types
  23. import Network.IRC.Util
  24. $(deriveLoggers "HSL" [HSL.INFO, HSL.ERROR])
  25. data RawIn = Line !UTCTime !Text | EOS deriving (Show, Eq)
  26. data In = Timeout | EOD | Msg !Message deriving (Show, Eq)
  27. formatCommand :: (Exception e) => BotConfig -> Message -> IO ([e], [Text])
  28. formatCommand botConfig@BotConfig { .. } message =
  29. map (second catMaybes . partitionEithers)
  30. . forM (defaultCommandFormatter : cmdFormatters) $ \formatter ->
  31. try . evaluate $ formatter botConfig message
  32. parseLine :: (Exception e)
  33. => BotConfig -> UTCTime -> Text -> Map MessageParserId [MessagePart]
  34. -> IO ([e], [Message], Map MessageParserId [MessagePart])
  35. parseLine botConfig@BotConfig { .. } time line msgParts =
  36. map mconcat . forM parsers $ \MessageParser { .. } -> do
  37. let parserMsgParts = concat . maybeToList $ lookup msgParserId msgParts
  38. let parserMsgPartsMap = singletonMap msgParserId parserMsgParts
  39. eresult <- try . evaluate $ msgParser botConfig time line parserMsgParts
  40. return $ case eresult of
  41. Left e -> ([e], [] , parserMsgPartsMap)
  42. Right ParseReject -> ([] , [] , parserMsgPartsMap)
  43. Right (ParsePartial msgParts') -> ([] , [] , singletonMap msgParserId msgParts')
  44. Right (ParseDone message msgParts') -> ([] , [message], singletonMap msgParserId msgParts')
  45. where
  46. parsers = defaultParsers ++ msgParsers
  47. sendCommandLoop :: MessageChannel Message -> Bot -> IO ()
  48. sendCommandLoop commandChan bot@Bot { .. } = do
  49. msg@(Message _ _ cmd) <- receiveMessage commandChan
  50. (exs, lines_) <- formatCommand botConfig msg
  51. forM_ exs $ \(ex :: SomeException) ->
  52. errorM ("Error while formatting command: " ++ show cmd ++ "\nError: " ++ show ex)
  53. unless (null lines_) $
  54. handle (\(e :: SomeException) -> do
  55. errorM ("Error while writing to connection: " ++ show e)
  56. closeMessageChannel commandChan) $
  57. forM_ lines_ $ \line -> do
  58. TF.hprint botSocket "{}\r\n" $ TF.Only line
  59. infoM . unpack $ "> " ++ line
  60. commandChanClosed <- isClosedMessageChannel commandChan
  61. unless commandChanClosed $
  62. case fromMessage cmd of
  63. Just QuitCmd -> closeMessageChannel commandChan
  64. _ -> sendCommandLoop commandChan bot
  65. readMessageLoop :: MVar BotStatus -> MessageChannel In -> Bot -> Int -> IO ()
  66. readMessageLoop mvBotStatus inChan Bot { .. } timeoutDelay = evalStateT loop mempty
  67. where
  68. msgPartTimeout = 10
  69. loop = do
  70. msgParts <- get
  71. botStatus <- readMVar mvBotStatus
  72. case botStatus of
  73. Disconnected -> io $ closeMessageChannel inChan
  74. _ -> do
  75. msgParts' <- io $ do
  76. mLine <- try $ timeout timeoutDelay readLine
  77. case mLine of
  78. Left (e :: SomeException) -> do
  79. errorM $ "Error while reading from connection: " ++ show e
  80. sendMessage inChan EOD >> return msgParts
  81. Right Nothing -> sendMessage inChan Timeout >> return msgParts
  82. Right (Just (Line time line)) -> do
  83. (exs, msgs, msgParts') <- parseLine botConfig time line msgParts
  84. forM_ exs $ \(ex :: SomeException) ->
  85. errorM ("Error while parsing line: " ++ unpack line ++ "\nError: " ++ show ex)
  86. forM_ msgs $ sendMessage inChan . Msg
  87. return msgParts'
  88. Right (Just EOS) -> sendMessage inChan EOD >> return msgParts
  89. limit <- io $ map (addUTCTime (- msgPartTimeout)) getCurrentTime
  90. put $ validMsgParts limit msgParts'
  91. loop
  92. where
  93. validMsgParts limit =
  94. foldl' (\m (k, v) -> insertWith (++) k [v] m) mempty
  95. . concat
  96. . filter ((> limit) . msgPartTime . snd . headEx . sortBy (flip $ comparing (msgPartTime . snd)))
  97. . groupAllOn (fst &&& msgPartTarget . snd)
  98. . asList . concatMap (uncurry (map . (,))) . mapToList
  99. readLine = do
  100. eof <- hIsEOF botSocket
  101. if eof
  102. then return EOS
  103. else mask $ \unmask -> do
  104. line <- map initEx . unmask $ hGetLine botSocket
  105. infoM . unpack $ "< " ++ line
  106. now <- getCurrentTime
  107. return $ Line now line
  108. messageProcessLoop :: MessageChannel In -> MessageChannel Message -> IRC ()
  109. messageProcessLoop inChan messageChan = loop 0
  110. where
  111. loop !idleFor = do
  112. status <- get
  113. Bot { .. } <- ask
  114. let nick = botNick botConfig
  115. let origNick = botOrigNick botConfig
  116. let mpass = CF.lookup "password" (config botConfig)
  117. nStatus <- io . mask_ $
  118. if idleFor >= (oneSec * botTimeout botConfig)
  119. then infoM "Timeout" >> return Disconnected
  120. else do
  121. when (status == Kicked) $
  122. threadDelay (5 * oneSec) >> newMessage JoinCmd >>= sendMessage messageChan
  123. mIn <- receiveMessage inChan
  124. case mIn of
  125. Timeout -> do
  126. idleMsg <- newMessage IdleMsg
  127. sendMessage messageChan idleMsg
  128. sendWhoisMessage nick origNick
  129. return Idle
  130. EOD -> infoM "Connection closed" >> return Disconnected
  131. Msg (msg@Message { .. }) -> do
  132. nStatus <- handleMsg nick origNick message mpass
  133. sendMessage messageChan msg
  134. return nStatus
  135. put nStatus
  136. case nStatus of
  137. Idle -> loop (idleFor + oneSec)
  138. Disconnected -> return ()
  139. NickNotAvailable -> return ()
  140. NickAvailable -> return ()
  141. _ -> loop 0
  142. where
  143. sendWhoisMessage nick origNick =
  144. when (nick /= origNick && idleFor /= 0 && idleFor `mod` (10 * oneSec) == 0) $
  145. (newMessage . WhoisCmd . nickToText $ origNick) >>= sendMessage messageChan
  146. handleMsg nick origNick message mpass
  147. | Just (JoinMsg user) <- fromMessage message, userNick user == nick =
  148. infoM "Joined" >> return Joined
  149. | Just (KickMsg { .. }) <- fromMessage message, kickedNick == nick =
  150. infoM "Kicked" >> return Kicked
  151. | Just NickInUseMsg <- fromMessage message =
  152. infoM "Nick already in use" >> return NickNotAvailable
  153. | Just (ModeMsg { .. }) <- fromMessage message, modeUser == Self = do
  154. whenJust mpass $ \pass -> do
  155. msg <- newMessage $ PrivMsgReply (User (Nick "NickServ") "") $ "IDENTIFY " ++ pass
  156. sendMessage messageChan msg
  157. newMessage JoinCmd >>= sendMessage messageChan
  158. return Connected
  159. | Just (WhoisNoSuchNick n) <- fromMessage message, n == origNick =
  160. infoM "Original nick available" >> return NickAvailable
  161. | otherwise =
  162. return Connected