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, resolvTimeout),
  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 import Timeout (Timeout(..))
  35
  36
  37 -- | Report results for this LookupResult. If there's a Nothing in the
  38 --   second component, there were no query results, so we report that
  39 --   the domain was not delegated. If there were some results and
  40 --   there are leftovers (after "cleaning"), we report those as well.
  41 report :: LookupResult -> IO ()
  42 report (d, Nothing) =
  43   putStrLn $ "Domain " ++ (show d) ++ " not delegated."
  44 report (d, Just leftovers) =
  45   unless (null leftovers) $
  46     putStrLn $ "Domain " ++
  47                (show d) ++
  48                " delegates somewhere else: " ++
  49                (join " " (map show leftovers))
  50
  51
  52 -- | "Clean" a lookup result by subtracting out the list of delegates.
  53 --   There shouldn't be anything left over. If there were no lookup
  54 --   results, we leave the Nothing in place so that 'report' can
  55 --   pattern match on it.
  56 clean :: [Domain] -- ^ List of delegates, @ds@
  57       -> LookupResult
  58       -> LookupResult
  59 clean _ p@(_, Nothing) = p
  60 clean delgts (d, Just targets) =
  61   (d, Just $ targets \\ delgts)
  62
  63
  64
  65 main :: IO ()
  66 main = do
  67   cfg <- get_args
  68
  69   -- This reads stdin.
  70   input <- BS.getContents
  71
  72   -- Split the input on any whitespace characters.
  73   let raw_domains = BS.words input
  74
  75   -- Convert these to ByteStrings.
  76   let raw_delegates = map BS.pack (delegates cfg)
  77
  78   let normalize_function =
  79         if (no_append_root cfg)
  80         then normalize_case
  81         else normalize
  82
  83   -- Normalize the given names and delegates
  84   let nrml_domains   = map normalize_function raw_domains
  85   let nrml_delegates = map normalize_function raw_delegates
  86
  87   rc' <- case (server cfg) of
  88            Nothing -> return defaultResolvConf
  89            Just s -> do
  90              s' <- resolve_address s
  91              case s' of
  92                Nothing -> do
  93                  hPutStrLn stderr ("Bad DNS server or lookup error: " ++ s)
  94                  exitWith (ExitFailure exit_bad_server)
  95                Just s'' ->
  96                  return $ defaultResolvConf { resolvInfo =
  97                                                 RCHostName (show s'') }
  98
  99   -- Set the timeout from the command line. The resolvTimeout field is
 100   -- in microseconds, so we multiply by one million.
 101   let rc = rc' { resolvTimeout = 1000 * 1000 * (seconds $ timeout cfg) }
 102   rs <- makeResolvSeed rc
 103
 104   let lookup_function = case cfg of
 105                           NS{} -> lookupNS'
 106                           MX{} -> lookupMX'
 107
 108   _ <- withResolver rs $ \resolver -> do
 109          -- Bad stuff happens if we try to run these lookups in parallel
 110          -- instead of the reports.
 111          records <- mapM (lookup_function resolver) nrml_domains
 112          let cleaned_records = map (clean nrml_delegates) records
 113          parallel (map report cleaned_records)
 114
 115   stopGlobalPool