src/Main.hs

   1 module Main
   2 where
   3
   4 import Control.Concurrent.ParallelIO.Global (
   5   parallelInterleaved,
   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   normalize,
  21   normalizeCase,
  22   withResolver )
  23 import System.Exit (ExitCode(..), exitWith)
  24 import System.IO (hPutStrLn, stderr)
  25
  26
  27 import CommandLine (Args(..), get_args)
  28 import DNS (
  29   LookupResult,
  30   lookupMX',
  31   lookupNS',
  32   resolve_address )
  33 import ExitCodes (exit_bad_server)
  34 import Timeout (Timeout(..))
  35
  36
  37 -- | Given a list of delegates, report results for this
  38 --   'LookupResult'.
  39 --
  40 --   If there's an empty list in the second component, there were no
  41 --   query results, so we report that the domain was not delegated. If
  42 --   there were some results and there are leftovers (after removing
  43 --   the delegates), we report those as well.
  44 --
  45 --   Before processing, all names are normalized using the supplied
  46 --   function @normalize_function@. Ideally this should be the same
  47 --   function applied to the user-input names.
  48 --
  49 report :: [Domain] -- ^ The list of @delgts@
  50        -> (Domain -> Domain) -- ^ Domain name normalization function,
  51                              --   @normalize_function@.
  52        -> LookupResult
  53        -> IO ()
  54
  55 -- If the lookup resulted in a DNS error, we just ignore the whole
  56 -- thing.
  57 report _ _ (_, Left _) = return ()
  58
  59 -- If the lookup succeeded but there were no results, report that the
  60 -- domain is not delegated.
  61 report _ _ (d, Right []) =
  62   putStrLn $ "Domain " ++ (show d) ++ " not delegated."
  63
  64 -- Otherwise, subtract our delegates from the list of results and
  65 -- report the leftovers.
  66 report delgts normalize_function (d, Right raw_hosts) = do
  67   let nrml_hosts = map normalize_function raw_hosts
  68   let leftovers = nrml_hosts \\ delgts
  69   unless (null leftovers) $
  70     putStrLn $ "Domain " ++
  71                (show d) ++
  72                " delegates somewhere else: " ++
  73                (join " " (map show leftovers))
  74
  75
  76 main :: IO ()
  77 main = do
  78   cfg <- get_args
  79
  80   -- This reads stdin.
  81   input <- BS.getContents
  82
  83   -- Split the input on any whitespace characters.
  84   let raw_domains = BS.words input
  85
  86   -- Convert these to ByteStrings.
  87   let raw_delegates = map BS.pack (delegates cfg)
  88
  89   let normalize_function =
  90         if (no_append_root cfg)
  91         then normalizeCase
  92         else normalize
  93
  94   -- Normalize the given names and delegates
  95   let nrml_domains   = map normalize_function raw_domains
  96   let nrml_delegates = map normalize_function raw_delegates
  97
  98   rc' <- case (server cfg) of
  99            Nothing -> return defaultResolvConf
 100            Just s -> do
 101              s' <- resolve_address s
 102              case s' of
 103                Left err -> do
 104                  let errmsg = show err
 105                  hPutStrLn stderr ("Bad DNS server or lookup error: " ++ errmsg)
 106                  exitWith (ExitFailure exit_bad_server)
 107                Right []  -> do
 108                  hPutStrLn stderr ("Hostname " ++ s ++ " has no 'A' records.")
 109                  exitWith (ExitFailure exit_bad_server)
 110                Right (srv:_) ->
 111                  return $ defaultResolvConf { resolvInfo =
 112                                                 RCHostName (show srv) }
 113
 114   -- Set the timeout from the command line. The resolvTimeout field is
 115   -- in microseconds, so we multiply by one million.
 116   let a_million = 1000 * 1000
 117   let rc = rc' { resolvTimeout = a_million * seconds (timeout cfg) }
 118   rs <- makeResolvSeed rc
 119
 120   let lookup_function = case cfg of
 121                           NS{} -> lookupNS'
 122                           MX{} -> lookupMX'
 123
 124   let lookup' d = withResolver rs $ \resolver ->
 125                     lookup_function resolver d
 126
 127   -- Construct a list of [IO whatever]. The withResolver calls
 128   -- are the ones that should be run in parallel.
 129   let actions = map lookup' nrml_domains
 130
 131   -- Run the lookup actions in parallel.
 132   results <- parallelInterleaved actions
 133
 134   -- Output the results.
 135   _ <- mapM (report nrml_delegates normalize_function) results
 136
 137   stopGlobalPool