packages feed

esotericbot-0.0.1: Esotericbot/Esotericbot.hs

{-# OPTIONS -XMagicHash #-}
import Prelude as P

import System.Environment
import qualified System.IO as IO
import System.Posix.Resource
import System.Posix.User
import System.Directory
import System.Posix.Env

import Network

import Data.ByteString.Unsafe
import Data.ByteString.Char8 as BSS
import Data.ByteString.Lazy.Char8 as BS
import Data.List.Stream as L
import Data.Char
import Data.Maybe
import qualified Data.Map as M
import Data.Monoid
import Data.Foldable as F

import Control.Applicative
import Control.Concurrent
import Control.Monad.Stream as C
import Control.Concurrent.STM

import Esotericbot.Execution
import Esotericbot.IRCCom
import Esotericbot.BSUtils
import Esotericbot.EBTypes as Ebt
import Esotericbot.Config
import Esotericbot.Safely ( both_limits )
import Esotericbot.Chroot as Ch
import Esotericbot.BSH

main = withSocketsDo $ do
   args <- getArgs
   if args == [ "help" ]
      then
         usage
      else do
         let mfp = get_conf_file args
             fp = maybe "esotericbot.conf" id mfp
         conf <- read_conf fp
         print conf
         set_mem conf
         ( vC , go_var ) <- liftIO $ atomically $ do mem <- newTVar 0; go <- newTVar True; return ( mem , go )
         flip evalStateT ( conf { vCurrent = vC , go = go_var } ) $ do -- stick the new info in the reader
            h <- Main.join
            -- set the user id, or chroot then set the user id, or do nothing if theres no chroot or userid
            F.sequence_  $ ( do ch <- Ebt.chroot conf
                                maybe ( error "You supply a chroot but not a user to setUID to: this is insecure." ) 
                                      ( \ u -> 
                                         Just $ liftIO $ do
                                            setCurrentDirectory ch
                                            Ch.chroot ch
                                            setUserID u
                                            setEnv "LD_LIBRARY_PATH" "." True )
                                      $ user conf )
                           <|> ( do u <- user conf
                                    Just $ liftIO $ setUserID u )
            state_fork $ do
               cmds <- listen h
               find_server_then_run_cmds h cmds 
            state_fork $ communicate h
            ping_every_minute h

ping_every_minute :: IO.Handle -> SmallBotM ( )
ping_every_minute h = do
   sb <- get
   ping <- liftIO $ C.liftM ( flip hAppend $ nick sb ) $ ls2bs 6 "PING :"#
   pinger
   if_go $ do
      writeMsg h ping
      ping_every_minute h 
   where
   pinger = 
     C.replicateM_ 60 $ if_go $ liftIO $ threadDelay 1000000
   

if_go f = do
   sb <- get
   b <- liftIO $ atomically $ readTVar $ go sb
   if b 
     then f
     else return $ error "borkborkbork"

set_mem =
   maybe ( return ( ) )
         ( setResourceLimit ResourceTotalMemory . both_limits )
         . mmem_limit

state_fork f = do
   r <- get
   liftIO $ forkIO $ evalStateT f r
   return ( )

get_conf_file args =
   case args of
      [ fname ] ->
         Just fname   
      _ ->
         Nothing

usage = do
   name <- getProgName
   IO.putStrLn $ usage_string name 

usage_string name = L.unlines 
   [ "Usage:"
   , name L.++ " ( [configuration_file] | help )"
   , "Esotericbot optionally takes one command line argument, which is either the string 'help', which will give you this information, or it will be the location of a configuration file."
   , "If this argument is not provided, esotericbot will look for a file called esotericbot.conf in the present working directory, and use that as the configuration file."
   ]

communicate :: IO.Handle -> SmallBotM ( )
communicate h = do
   sb <-get
   if_go $ do
      raw_msg <- liftIO $ BSS.getLine
      quit <- liftIO $ unsafePackAddressLen 4 "QUIT"#
      let should_stop2 = raw_msg == quit
      possible_failure should_stop2
      let mmsg = unescape raw_msg
      maybe ( liftIO $ IO.putStrLn "Input error: '\\' can only be followed by an integer, or another '\\'." )
            ( writeMsg h )
            mmsg
      if not should_stop2
         then communicate h 
         else return ( )
   where
   unescape bs =
      if BSS.elem '\\' bs
         then 
            let bss = BSS.split '\\' bs
                fbs = L.head bss
                rbs = L.tail bss
                unescaped = 
                  L.map ( \ b ->
                           let mintb = BSS.readInt b in
                           maybe ( let mcr = BSS.uncons b in
                                   maybe Nothing
                                         ( \ ( c , r ) -> 
                                            if c == '\\' 
                                               then Just $ '\\' `BSS.cons` r
                                               else Nothing
                                         ) mcr
                                 )
                                 ( \ ( i , r ) ->
                                    Just $ chr i `BSS.cons` r
                                 )
                                 mintb 
                        ) rbs
                in C.sequence unescaped >>= Just . BSS.concat >>= Just . BSS.append fbs
         else Just bs 
   possible_failure b = do
      sb <- get
      if b 
         then liftIO $ atomically $ writeTVar ( go sb ) False 
         else return ( )


port   = 6667

join :: SmallBotM IO.Handle
join = do
   sb <- get
   h <- liftIO $ connectTo ( server sb ) $ PortNumber $ fromIntegral port
   liftIO $ IO.hSetBuffering h IO.NoBuffering
   n <- liftIO $ ls2bs 4 "NICK"# 
   u <- liftIO $ ls2bs 5 "USER "# 
   j <- liftIO $ ls2bs 4 "JOIN"#
   write2 h n $ nick sb
   n <- liftIO getProgName
   eso <- liftIO $ ls2bs 7 " 0 * : "#
   writeMsg h $ u `hAppend` nick sb `hAppend` eso `hAppend` BSHString n
   C.mapM_ ( write2 h j ) $ chans sb -- Join these channels
   let password = pass sb
   if BS.null password
      then return ( )
      else do
         p <- liftIO $ ls2bs 27 "PRIVMSG NickServ :identify "#
         writeMsg h $ p `hAppend` password 
   return h

listen h = do
   sb <- get
   liftIO $ 
      catch ( do s <- BS.hGetContents h
                 let lines = BS.split '\n' s
                 return lines
            )
            $ const $ do liftIO $ do IO.putStrLn "The handle broke.  Buy a better door."
                                     IO.hClose h
                                     atomically $ writeTVar ( go sb ) False
                         error "Borkborkbork"

-- Nicked from Data.Foldable, but using Data.List.Stream
foldlM :: Monad m => (a -> b -> m a) -> a -> [ b ] -> m a
foldlM f z0 xs = L.foldr f' return xs z0
  where f' x k z = f z x >>= k

find_server_then_run_cmds h cmds = do
   rest <- find_connected_server cmds
   help_str <- liftIO $ ls2bs 4 "help"#
   op_str <- liftIO $ ls2bs 2 "op"#
   sb <- get
   s <- Main.foldlM ( \ m c -> 
                       fmap ( flip ( M.insert c ) m ) 
                            $ liftIO $ atomically 
                                        $ newTVar initial_op_search 
               ) M.empty $ chans sb
   l <- Main.foldlM ( \ m c -> 
                       fmap ( flip ( M.insert c ) m )
                            $ liftIO $ atomically
                                        $ newTVar True
               ) M.empty $ chans sb
   put $ sb { searching = s , listening = l } 
   C.mapM_ ( run_cmd h help_str op_str )  rest

find_connected_server all_cmds = do
   let cmd = L.head all_cmds
       cmds = L.tail all_cmds
   sb <- get
   if_go $ do
      liftIO $ do BS.putStrLn cmd 
                  IO.hFlush IO.stdout
      pres <- get_real_server cmd
      either ( const $ find_connected_server cmds )
             ( \ cs -> do put $ sb { cserver = Just cs } 
                          return cmds )
             $ snd $ pres

-- if the line is an operator command, we have to multiplex here while waiting for the op names
-- otherwise, it could be a help request
-- if it's not that, then try to run a plugin

run_cmd h help_str op_str line = do
   sb <- get
   should_go <- liftIO $ atomically $ readTVar $ go sb
   if should_go
      then do
         liftIO $ do BS.putStrLn line
                     IO.hFlush IO.stdout
         any_underway <- liftIO $ atomically $ C.liftM L.or $ C.mapM ( \ v -> readTVar v >>= return . underway ) $ M.elems $ searching sb
         if any_underway
            then do
               mend <- end_of_names line
               maybe ( do mcops <- get_ops line
                          maybe smallbot_command
                                add_ops
                                mcops
                     )
                     end_ops_search 
                     mend
            else smallbot_command
       else
         return ( )
   where
   add_ops ( chan , new_ops ) = do
      sb <- get
      liftIO $ atomically $ do
         let Just op_search_var = M.lookup chan $ searching sb 
         op_search <- readTVar op_search_var
         writeTVar op_search_var $ op_search { ops = new_ops L.++ ops op_search }
   end_ops_search chan = do
      sb <- get
      liftIO $ atomically $ do
         let Just op_search_var = M.lookup chan $ searching sb
         op_search <- readTVar op_search_var
         writeTVar op_search_var $ op_search { done = True  }
   smallbot_command = state_fork $ do -- wrapper for smallbot commands
      eb <- get
      if L.or $ L.map ( flip BS.elem line ) $ ( BS.head $ command_prefix eb ) : ( L.map fromJust $ L.filter isJust $ L.map ( \ l -> shortcut l >>= return . BS.head ) $ plugins eb )   -- a quick check!
         then sb_message
         else
            return ( )
   sb_message = do -- a smallbot command
      sb <- get
      parse_res <- cmd_line line
      either ( const $ return ( ) )
             ( either ( flip run_command $ UnknownPluginContext h help_str op_str )
                      ( flip run_command $ KnownPluginContext h )
             )
             parse_res