Add command-line processing.

[email-validator.git] / src / Main.hs

diff --git a/src/Main.hs b/src/Main.hs
index bd3cc45f5ff741e936243c9de0cf5baa84f0e1ce..f38a7db4b3f5c1d60a2067c34d18c2af512e29ec 100644 (file)
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,9 +1,11 @@
 {-# LANGUAGE DoAndIfThenElse #-}
+{-# LANGUAGE RecordWildCards #-}
 
 module Main
 where
 
 import Control.Concurrent.ParallelIO.Global (parallel, stopGlobalPool)
+import Control.Monad (when)
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.UTF8 as BSU
 import Network.DNS (
@@ -14,9 +16,21 @@ import Network.DNS (
   makeResolvSeed,
   withResolver)
 import Network.DNS.Lookup (lookupMX)
-import System.IO (hGetContents, hFlush, hPutStrLn, stdin, stdout)
+import System.Directory (doesFileExist)
+import System.Exit (exitWith, ExitCode(..))
+import System.IO (
+  Handle,
+  IOMode( WriteMode ),
+  hClose,
+  hFlush,
+  openFile,
+  stdin,
+  stdout)
 import Text.Regex.PCRE.Light (compile, match, utf8)
 
+import CommandLine (Args(..), apply_args)
+import ExitCodes (exit_input_file_doesnt_exist)
+
 type Address = BSU.ByteString
 
 -- | Resolver parameters. We increase the default timeout from 3 to 5
@@ -96,19 +110,40 @@ validate resolver address = do
   else do
     return (address, False)
 
+
+append_handle_with_newline :: Handle -> BS.ByteString -> IO ()
+append_handle_with_newline h bs = do
+  BS.hPutStr h bs
+  BS.hPutStr h newline
+  where
+    newline = BSU.fromString "\n"
+
+
 main :: IO ()
 main = do
-  input <- hGetContents stdin
+  Args{..} <- apply_args
+
+  input <- case input_file of
+             Nothing   -> BS.hGetContents stdin
+             Just path -> do
+               is_file <- doesFileExist path
+               when (not is_file) $ do
+                 exitWith (ExitFailure exit_input_file_doesnt_exist)
+               BS.readFile path
 
-  let addresses = BSU.lines $ BSU.fromString input
+  output_handle <- case output_file of
+                     Nothing -> return stdout
+                     Just path -> openFile path WriteMode
+
+  let addresses = BSU.lines input
   let nonempty_addresses = filter (not . BS.null) addresses
 
   rs <- makeResolvSeed resolv_conf
   withResolver rs $ \resolver -> do
     let actions = map (validate resolver) nonempty_addresses
     results <- parallel actions
+    stopGlobalPool
     let good_pairs = filter snd results
-    mapM_ ((hPutStrLn stdout) . BSU.toString . fst) good_pairs
-
-  stopGlobalPool
-  hFlush stdout
+    mapM_ ((append_handle_with_newline output_handle) . fst) good_pairs
+    hFlush output_handle
+    hClose output_handle