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.4KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  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)
  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. forM_ lines_ $ \line -> do
  54. handle (\(e :: SomeException) -> do
  55. errorM ("Error while writing to connection: " ++ show e)
  56. closeMessageChannel commandChan) $ do
  57. TF.hprint botSocket "{}\r\n" $ TF.Only line
  58. infoM . unpack $ "> " ++ line
  59. commandChanClosed <- isClosedMessageChannel commandChan
  60. unless commandChanClosed $
  61. case fromMessage cmd of
  62. Just QuitCmd -> closeMessageChannel commandChan
  63. _ -> sendCommandLoop commandChan bot
  64. readMessageLoop :: MVar BotStatus -> MessageChannel In -> Bot -> Int -> IO ()
  65. readMessageLoop mvBotStatus inChan Bot { .. } timeoutDelay = loop mempty
  66. where
  67. msgPartTimeout = 10
  68. loop msgParts = do
  69. botStatus <- readMVar mvBotStatus
  70. case botStatus of
  71. Disconnected -> io $ closeMessageChannel inChan
  72. _ -> do
  73. msgParts' <- io $ do
  74. mLine <- try $ timeout timeoutDelay readLine
  75. case mLine of
  76. Left (e :: SomeException) -> do
  77. errorM $ "Error while reading from connection: " ++ show e
  78. sendMessage inChan EOD >> return msgParts
  79. Right Nothing -> sendMessage inChan Timeout >> return msgParts
  80. Right (Just (Line time line)) -> do
  81. (exs, msgs, msgParts') <- parseLine botConfig time line msgParts
  82. forM_ exs $ \(ex :: SomeException) ->
  83. errorM ("Error while parsing line: " ++ unpack line ++ "\nError: " ++ show ex)
  84. forM_ msgs $ sendMessage inChan . Msg
  85. return msgParts'
  86. Right (Just EOS) -> sendMessage inChan EOD >> return msgParts
  87. limit <- io $ map (addUTCTime (- msgPartTimeout)) getCurrentTime
  88. loop $ validMsgParts limit msgParts'
  89. validMsgParts limit =
  90. foldl' (\m (k, v) -> insertWith (++) k [v] m) mempty
  91. . concat
  92. . filter ((> limit) . msgPartTime . snd . headEx . sortBy (flip $ comparing (msgPartTime . snd)))
  93. . groupAllOn (fst &&& msgPartTarget . snd)
  94. . asList
  95. . concatMap (uncurry (map . (,)))
  96. . mapToList
  97. readLine = do
  98. eof <- hIsEOF botSocket
  99. if eof
  100. then return EOS
  101. else mask $ \unmask -> do
  102. line <- map initEx . unmask $ hGetLine botSocket
  103. infoM . unpack $ "< " ++ line
  104. now <- getCurrentTime
  105. return $ Line now line
  106. messageProcessLoop :: MessageChannel In -> MessageChannel Message -> IRC ()
  107. messageProcessLoop inChan messageChan = loop 0
  108. where
  109. loop !idleFor = do
  110. status <- get
  111. Bot { .. } <- ask
  112. let nick = botNick botConfig
  113. let origNick = botOrigNick botConfig
  114. let mpass = CF.lookup "password" (config botConfig)
  115. nStatus <- io . mask_ $
  116. if idleFor >= (oneSec * botTimeout botConfig)
  117. then infoM "Timeout" >> return Disconnected
  118. else do
  119. when (status == Kicked) $
  120. threadDelay (5 * oneSec) >> (sendMessage messageChan =<< newMessage JoinCmd)
  121. mIn <- receiveMessage inChan
  122. case mIn of
  123. Timeout -> do
  124. sendMessage messageChan =<< newMessage IdleMsg
  125. sendWhoisMessage nick origNick idleFor
  126. return Idle
  127. EOD -> infoM "Connection closed" >> return Disconnected
  128. Msg (msg@Message { .. }) -> do
  129. nStatus <- handleMsg nick origNick message mpass
  130. sendMessage messageChan msg
  131. return nStatus
  132. put nStatus
  133. case nStatus of
  134. Idle -> loop (idleFor + oneSec)
  135. Disconnected -> return ()
  136. NickNotAvailable -> return ()
  137. NickAvailable -> return ()
  138. _ -> loop 0
  139. sendWhoisMessage nick origNick idleFor =
  140. when (nick /= origNick && idleFor /= 0 && idleFor `mod` (10 * oneSec) == 0) $
  141. sendMessage messageChan =<< (newMessage . WhoisCmd . nickToText $ origNick)
  142. handleMsg nick origNick message mpass
  143. | Just (JoinMsg user) <- fromMessage message, userNick user == nick =
  144. infoM "Joined" >> return Joined
  145. | Just (KickMsg { .. }) <- fromMessage message, kickedNick == nick =
  146. infoM "Kicked" >> return Kicked
  147. | Just NickInUseMsg <- fromMessage message =
  148. infoM "Nick already in use" >> return NickNotAvailable
  149. | Just (ModeMsg { .. }) <- fromMessage message, modeUser == Self = do
  150. whenJust mpass $ \pass -> do
  151. msg <- newMessage $ PrivMsgReply (User (Nick "NickServ") "") $ "IDENTIFY " ++ pass
  152. sendMessage messageChan msg
  153. sendMessage messageChan =<< newMessage JoinCmd
  154. return Connected
  155. | Just (WhoisNoSuchNickMsg n) <- fromMessage message, n == origNick =
  156. infoM "Original nick available" >> return NickAvailable
  157. | otherwise =
  158. return Connected