module Main where import Control.Concurrent.ParallelIO.Global ( parallel, stopGlobalPool ) import Control.Monad (unless) import qualified Data.ByteString.Char8 as BS ( getContents, pack, words ) import Data.List ((\\)) import Data.String.Utils (join) import Network.DNS ( Domain, FileOrNumericHost(RCHostName), ResolvConf(resolvInfo, resolvTimeout), defaultResolvConf, makeResolvSeed, withResolver ) import System.Exit (ExitCode(..), exitWith) import System.IO (hPutStrLn, stderr) import CommandLine (Args(..), get_args) import DNS ( LookupResult, lookupMX', lookupNS', normalize, normalize_case, resolve_address ) import ExitCodes (exit_bad_server) import Timeout (Timeout(..)) -- | Report results for this LookupResult. If there's a Nothing in the -- second component, there were no query results, so we report that -- the domain was not delegated. If there were some results and -- there are leftovers (after "cleaning"), we report those as well. report :: LookupResult -> IO () report (d, Nothing) = putStrLn $ "Domain " ++ (show d) ++ " not delegated." report (d, Just leftovers) = unless (null leftovers) $ putStrLn $ "Domain " ++ (show d) ++ " delegates somewhere else: " ++ (join " " (map show leftovers)) -- | "Clean" a lookup result by subtracting out the list of delegates. -- There shouldn't be anything left over. If there were no lookup -- results, we leave the Nothing in place so that 'report' can -- pattern match on it. clean :: [Domain] -- ^ List of delegates, @ds@ -> LookupResult -> LookupResult clean _ p@(_, Nothing) = p clean delgts (d, Just targets) = (d, Just $ targets \\ delgts) main :: IO () main = do cfg <- get_args -- This reads stdin. input <- BS.getContents -- Split the input on any whitespace characters. let raw_domains = BS.words input -- Convert these to ByteStrings. let raw_delegates = map BS.pack (delegates cfg) let normalize_function = if (no_append_root cfg) then normalize_case else normalize -- Normalize the given names and delegates let nrml_domains = map normalize_function raw_domains let nrml_delegates = map normalize_function raw_delegates rc' <- case (server cfg) of Nothing -> return defaultResolvConf Just s -> do s' <- resolve_address s case s' of Nothing -> do hPutStrLn stderr ("Bad DNS server or lookup error: " ++ s) exitWith (ExitFailure exit_bad_server) Just s'' -> return $ defaultResolvConf { resolvInfo = RCHostName (show s'') } -- Set the timeout from the command line. The resolvTimeout field is -- in microseconds, so we multiply by one million. let rc = rc' { resolvTimeout = 1000 * 1000 * (seconds $ timeout cfg) } rs <- makeResolvSeed rc let lookup_function = case cfg of NS{} -> lookupNS' MX{} -> lookupMX' _ <- withResolver rs $ \resolver -> do -- Bad stuff happens if we try to run these lookups in parallel -- instead of the reports. records <- mapM (lookup_function resolver) nrml_domains let cleaned_records = map (clean nrml_delegates) records parallel (map report cleaned_records) stopGlobalPool