src/Main.hs

   1 module Main
   2 where
   3
   4 import Control.Concurrent.ParallelIO.Global (
   5   parallel,
   6   stopGlobalPool )
   7 import Control.Monad (unless)
   8 import qualified Data.ByteString.Char8 as BS (
   9   getContents,
  10   pack,
  11   words )
  12 import Data.List ((\\))
  13 import Data.String.Utils (join)
  14 import Network.DNS (
  15   Domain,
  16   FileOrNumericHost(RCHostName),
  17   ResolvConf(resolvInfo),
  18   defaultResolvConf,
  19   makeResolvSeed,
  20   withResolver )
  21 import System.Exit (ExitCode(..), exitWith)
  22 import System.IO (hPutStrLn, stderr)
  23
  24
  25 import CommandLine (Args(..), get_args)
  26 import DNS (
  27   LookupResult,
  28   lookupMX',
  29   lookupNS',
  30   normalize,
  31   normalize_case,
  32   resolve_address )
  33 import ExitCodes (exit_bad_server)
  34
  35
  36 -- | Report results for this LookupResult. If there's a Nothing in the
  37 --   second component, there were no query results, so we report that
  38 --   the domain was not delegated. If there were some results and
  39 --   there are leftovers (after "cleaning"), we report those as well.
  40 report :: LookupResult -> IO ()
  41 report (d, Nothing) =
  42   putStrLn $ "Domain " ++ (show d) ++ " not delegated."
  43 report (d, Just leftovers) =
  44   unless (null leftovers) $
  45     putStrLn $ "Domain " ++
  46                (show d) ++
  47                " delegates somewhere else: " ++
  48                (join " " (map show leftovers))
  49
  50
  51 -- | "Clean" a lookup result by subtracting out the list of delegates.
  52 --   There shouldn't be anything left over. If there were no lookup
  53 --   results, we leave the Nothing in place so that 'report' can
  54 --   pattern match on it.
  55 clean :: [Domain] -- ^ List of delegates, @ds@
  56       -> LookupResult
  57       -> LookupResult
  58 clean _ p@(_, Nothing) = p
  59 clean delgts (d, Just targets) =
  60   (d, Just $ targets \\ delgts)
  61
  62
  63
  64 main :: IO ()
  65 main = do
  66   cfg <- get_args
  67
  68   -- This reads stdin.
  69   input <- BS.getContents
  70
  71   -- Split the input on any whitespace characters.
  72   let raw_domains = BS.words input
  73
  74   -- Convert these to ByteStrings.
  75   let raw_delegates = map BS.pack (delegates cfg)
  76
  77   let normalize_function =
  78         if (no_append_root cfg)
  79         then normalize_case
  80         else normalize
  81
  82   -- Normalize the given names and delegates
  83   let nrml_domains   = map normalize_function raw_domains
  84   let nrml_delegates = map normalize_function raw_delegates
  85
  86   rc <- case (server cfg) of
  87           Nothing -> return defaultResolvConf
  88           Just s -> do
  89             s' <- resolve_address s
  90             case s' of
  91               Nothing -> do
  92                 hPutStrLn stderr ("Bad DNS server or lookup error: " ++ s)
  93                 exitWith (ExitFailure exit_bad_server)
  94               Just s'' ->
  95                 return $ defaultResolvConf { resolvInfo =
  96                                                RCHostName (show s'') }
  97
  98   rs <- makeResolvSeed rc
  99
 100   let lookup_function = case cfg of
 101                           NS{} -> lookupNS'
 102                           MX{} -> lookupMX'
 103
 104   _ <- withResolver rs $ \resolver -> do
 105          -- Bad stuff happens if we try to run these lookups in parallel
 106          -- instead of the reports.
 107          records <- mapM (lookup_function resolver) nrml_domains
 108          let cleaned_records = map (clean nrml_delegates) records
 109          parallel (map report cleaned_records)
 110
 111   stopGlobalPool