module Main where import Control.Concurrent.ParallelIO.Global ( parallelInterleaved, 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, normalize, normalizeCase, withResolver ) import System.Exit (ExitCode(ExitFailure), exitWith) import System.IO (hPutStrLn, stderr) import CommandLine ( Args(NS,MX,delegates,no_append_root,server,timeout), get_args) import DNS ( LookupResult, lookupMX', lookupNS', resolve_address ) import ExitCodes (exit_bad_server) import Timeout (Timeout(seconds)) -- | Given a list of delegates, report results for this -- 'LookupResult'. -- -- If there's an empty list 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 removing -- the delegates), we report those as well. -- -- Before processing, all names are normalized using the supplied -- function @normalize_function@. Ideally this should be the same -- function applied to the user-input names. -- report :: [Domain] -- ^ The list of @delgts@ -> (Domain -> Domain) -- ^ Domain name normalization function, -- @normalize_function@. -> LookupResult -> IO () -- If the lookup resulted in a DNS error, we just ignore the whole -- thing. report _ _ (_, Left _) = return () -- If the lookup succeeded but there were no results, report that the -- domain is not delegated. report _ _ (d, Right []) = putStrLn $ "Domain " ++ (show d) ++ " not delegated." -- Otherwise, subtract our delegates from the list of results and -- report the leftovers. report delgts normalize_function (d, Right raw_hosts) = do let nrml_hosts = map normalize_function raw_hosts let leftovers = nrml_hosts \\ delgts unless (null leftovers) $ putStrLn $ "Domain " ++ (show d) ++ " delegates somewhere else: " ++ (join " " (map show leftovers)) 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 normalizeCase 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 Left err -> do let errmsg = show err hPutStrLn stderr ("Bad DNS server or lookup error: " ++ errmsg) exitWith (ExitFailure exit_bad_server) Right [] -> do hPutStrLn stderr ("Hostname " ++ s ++ " has no 'A' records.") exitWith (ExitFailure exit_bad_server) Right (srv:_) -> return $ defaultResolvConf { resolvInfo = RCHostName (show srv) } -- Set the timeout from the command line. The resolvTimeout field is -- in microseconds, so we multiply by one million. let a_million = 1000 * 1000 :: Int let rc = rc' { resolvTimeout = a_million * seconds (timeout cfg) } rs <- makeResolvSeed rc let lookup_function = case cfg of NS{} -> lookupNS' MX{} -> lookupMX' let lookup' d = withResolver rs $ \resolver -> lookup_function resolver d -- Construct a list of [IO whatever]. The withResolver calls -- are the ones that should be run in parallel. let actions = map lookup' nrml_domains -- Run the lookup actions in parallel. results <- parallelInterleaved actions -- Output the results. _ <- mapM (report nrml_delegates normalize_function) results stopGlobalPool