src/Main.hs

   1 {-# LANGUAGE BangPatterns #-}
   2 {-# LANGUAGE DoAndIfThenElse #-}
   3
   4 module Main
   5 where
   6
   7 import Control.Concurrent (threadDelay)
   8 import Control.Exception.Base (bracket)
   9 import Control.Monad (when)
  10 import Data.List (isPrefixOf)
  11 import Data.Maybe (isNothing)
  12 import Data.Monoid ((<>))
  13 import Network (
  14   connectTo,
  15   PortID (PortNumber) )
  16 import System.Console.CmdArgs (def)
  17 import System.Directory (doesFileExist)
  18 import System.Exit (ExitCode(..), exitWith)
  19 import System.FilePath ((</>))
  20 import System.IO (
  21   BufferMode (NoBuffering),
  22   Handle,
  23   hClose,
  24   hGetChar,
  25   hGetLine,
  26   hPutStr,
  27   hSetBuffering,
  28   stderr,
  29   stdout )
  30 import System.IO.Error (catchIOError)
  31 import System.Log.Logger ( getLogger, rootLoggerName, saveGlobalLogger )
  32 import System.Timeout (timeout)
  33
  34 import CommandLine (get_args)
  35 import Configuration (Configuration(..), merge_optional)
  36 import ExitCodes (
  37   exit_no_feed_hosts,
  38   exit_no_password,
  39   exit_no_username )
  40 import Logging (
  41   init_logging,
  42   log_debug,
  43   log_error,
  44   log_info,
  45   log_warning )
  46 import qualified OptionalConfiguration as OC (
  47   OptionalConfiguration(..),
  48   from_rc )
  49 import Terminal (putGreenLn)
  50 import TSN.FeedHosts (FeedHosts(..))
  51 import TSN.Xml (parse_xmlfid, xml_prologue)
  52
  53
  54 -- | Receive a single line of text from a Handle, and send it to the
  55 --   debug log.
  56 --
  57 recv_line :: Handle -> IO String
  58 recv_line h = do
  59   line <- hGetLine h
  60   log_debug line
  61   return line
  62
  63
  64 -- | Takes a Configuration, and an XML document (as a String). The XML
  65 --   document is written to the output directory, as specified by the
  66 --   Configuration.
  67 --
  68 --   This can fail, but we don't purposefully throw any exceptions. If
  69 --   something goes wrong, we would rather log it and keep going.
  70 --
  71 save_document :: Configuration -> String -> IO ()
  72 save_document cfg doc =
  73   case maybe_path of
  74     Nothing ->
  75       log_error "Document missing XML_File_ID element."
  76     Just path -> do
  77       already_exists <- doesFileExist path
  78       when already_exists $ do
  79         let msg = "File " ++ path ++ " already exists, overwriting."
  80         log_warning msg
  81       writeFile path doc
  82       log_info $ "Wrote file: " ++ path ++ "."
  83   where
  84     xmlfid = fmap show (parse_xmlfid doc)
  85     filename = fmap (++ ".xml") xmlfid
  86     maybe_path = fmap ((output_directory cfg) </>) filename
  87
  88
  89 -- | Loop forever, writing the buffer to file whenever a new XML
  90 --   prologue is seen. This is the low-level "loop forever" function
  91 --   that we stay in as long as we are connected to one feed.
  92 --
  93 loop :: Configuration -> Handle -> [String] -> IO ()
  94 loop !cfg !h !buffer = do
  95   line <- recv_line h
  96
  97   if (xml_prologue `isPrefixOf` line && not (null buffer))
  98   then do
  99     -- This is the beginning of a new document, and we have an "old"
 100     -- one to save. The buffer is in reverse (newest first) order,
 101     -- though, so we have to reverse it first. We then concatenate all
 102     -- of its lines into one big string.
 103     let document = concat $ reverse buffer
 104     save_document cfg document
 105     loop cfg h [line] -- empty the buffer before looping again
 106   else
 107     -- append line to the head of the buffer and loop
 108     loop cfg h (line : buffer)
 109
 110
 111 log_in :: Configuration -> Handle -> IO ()
 112 log_in cfg h = do
 113   prompt1 <- recv_prompt h
 114
 115   if prompt1 /= username_prompt then
 116     log_error "Didn't receive username prompt."
 117   else do
 118     send_line h (username cfg)
 119     prompt2 <- recv_prompt h
 120
 121     if prompt2 /= password_prompt then
 122       log_error "Didn't receive password prompt."
 123     else do
 124       send_line h (password cfg)
 125       _ <- recv_line h -- "The Sports Network"
 126       return ()
 127   where
 128     username_prompt = "Username: "
 129     password_prompt = "Password: "
 130
 131     send_line :: Handle -> String -> IO ()
 132     send_line h' s = do
 133       hPutStr h' (s ++ "\r\n")
 134       putGreenLn s
 135
 136     recv_chars :: Int -> Handle -> IO String
 137     recv_chars n h' = do
 138       s <- sequence [ hGetChar h' | _ <- [1..n] ]
 139       putStr s
 140       return s
 141
 142     recv_prompt :: Handle -> IO String
 143     recv_prompt = recv_chars 10
 144
 145 connect_and_loop :: Configuration -> String -> IO ()
 146 connect_and_loop cfg host = do
 147   log_info $ "Connecting to " ++ host ++ "..."
 148   bracket acquire_handle release_handle action
 149   return ()
 150   where
 151     five_seconds :: Int
 152     five_seconds = 5000000
 153
 154     acquire_handle = connectTo host (PortNumber 4500)
 155     release_handle = hClose
 156     action h = do
 157       -- No buffering anywhere.
 158       hSetBuffering h NoBuffering
 159
 160       -- The feed is often unresponsive after we send out username. It
 161       -- happens in a telnet session, too (albeit less frequently?),
 162       -- so there might be a bug on their end.
 163       --
 164       -- If we dump the packets with tcpdump, it looks like their
 165       -- software is getting confused: they send us some XML in
 166       -- the middle of the log-in procedure. In any case, the easiest
 167       -- fix is to disconnect and try again.
 168       --
 169       login_worked <- timeout five_seconds $ log_in cfg h
 170       case login_worked of
 171         Nothing -> log_info "Login timed out (5s)."
 172         Just _ ->  loop cfg h []
 173
 174
 175 -- | A wrapper around threadDelay which takes seconds instead of
 176 --   microseconds as its argument.
 177 --
 178 thread_sleep :: Int -> IO ()
 179 thread_sleep seconds = do
 180   let microseconds = seconds * (10 ^ (6 :: Int))
 181   threadDelay microseconds
 182
 183
 184 -- | The entry point of the program.
 185 main :: IO ()
 186 main = do
 187   init_logging
 188   root_logger <- getLogger rootLoggerName
 189   saveGlobalLogger root_logger
 190
 191   rc_cfg <- OC.from_rc
 192   cmd_cfg <- get_args
 193
 194   -- Merge the config file options with the command-line ones,
 195   -- prefering the command-line ones.
 196   let opt_config = rc_cfg <> cmd_cfg
 197
 198   -- This is necessary because if the user specifies an empty list of
 199   -- hostnames in e.g. the config file, we want to bail rather than
 200   -- fall back on the default list (which gets merged from a
 201   -- Configuration below).
 202   when (null $ get_feed_hosts (OC.feed_hosts opt_config)) $ do
 203     log_error "No feed hosts supplied."
 204     exitWith (ExitFailure exit_no_feed_hosts)
 205
 206   when (isNothing (OC.password opt_config)) $ do
 207     log_error "No password supplied."
 208     exitWith (ExitFailure exit_no_password)
 209
 210   when (isNothing (OC.username opt_config)) $ do
 211     log_error "No username supplied."
 212     exitWith (ExitFailure exit_no_username)
 213
 214   -- Finally, update a default config with any options that have been
 215   -- set in either the config file or on the command-line.
 216   let cfg = (def :: Configuration) `merge_optional` opt_config
 217
 218   -- This may be superstition (and I believe stderr is unbuffered),
 219   -- but it can't hurt.
 220   hSetBuffering stderr NoBuffering
 221   hSetBuffering stdout NoBuffering
 222
 223   -- Begin connecting to our feed hosts, starting with the first one.
 224   round_robin cfg 0
 225
 226   where
 227     -- | This is the top-level "loop forever" function. If an
 228     --   exception is thrown, it will propagate up to this point, where
 229     --   it will be logged and ignored in style.
 230     --
 231     --   Afterwards, we recurse (call ourself) again to loop more forevers.
 232     --
 233     round_robin :: Configuration -> Int -> IO ()
 234     round_robin cfg feed_host_idx = do
 235       let hosts = get_feed_hosts $ feed_hosts cfg
 236       let host = hosts !! feed_host_idx
 237       catchIOError (connect_and_loop cfg host) (log_error . show)
 238       thread_sleep 10 -- Wait 10s before attempting to reconnect.
 239       round_robin cfg $ (feed_host_idx + 1) `mod` (length hosts)