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