module DNS (
  LookupResult,
  lookupMX',
  lookupNS',
  resolve_address )
where

import qualified Data.ByteString.Char8 as BS ( pack )
import Data.IP (IPv4)
import Network.DNS (
  Domain,
  DNSError,
  Resolver,
  defaultResolvConf,
  lookupA,
  lookupMX,
  lookupNS,
  lookupNSAuth,
  makeResolvSeed,
  withResolver )
import Text.Read ( readMaybe )

type LookupResult = (Domain, Either DNSError [Domain])

-- | Takes a String representing either a hostname or an IP
--   address. If a hostname was supplied, it is resolved to either an
--   [IPv4] or an error. If an IP address is supplied, it is returned
--   as a singleton [IPv4].
--
--   Examples:
--
--   >>> resolve_address "example.com"
--   Right [93.184.215.14]
--   >>> resolve_address "93.184.216.34"
--   Right [93.184.216.34]
--
resolve_address :: String -> IO (Either DNSError [IPv4])
resolve_address s =
  case read_result of
    Just addr  -> return $ Right [addr]
    Nothing -> do
      default_rs <- makeResolvSeed defaultResolvConf
      withResolver default_rs $ \resolver ->
        lookupA resolver (BS.pack s)
  where
    read_result :: Maybe IPv4
    read_result = readMaybe s


-- | This calls lookupMX, and returns the result as the second
--   component of a tuple whose first component is the domain name
--   itself.
--
--   Examples:
--
--   The example domain, example.com, has a NULLMX record.
--
--   >>> rs <- makeResolvSeed defaultResolvConf
--   >>> let domain = BS.pack "example.com."
--   >>> withResolver rs $ \resolver -> lookupMX' resolver domain
--   ("example.com.",Right ["."])
--
lookupMX' :: Resolver -> Domain -> IO LookupResult
lookupMX' resolver domain =
  fmap (pair_em . drop_priority) $ lookupMX resolver domain
  where
    drop_priority :: Either DNSError [(Domain, Int)]
                  -> Either DNSError [Domain]
    drop_priority = fmap (map fst)

    pair_em :: a -> (Domain, a)
    pair_em = (,) domain


-- | This calls lookupNS, and returns the result as the second
--   component of a tuple whose first component is the domain name
--   itself.
--
--   Examples:
--
--   The example domain, example.com, does have NS records, but the
--   order in which they are returned is variable, so we have to sort
--   them to get a reliable result.
--
--   >>> import Data.List (sort)
--   >>> import Control.Applicative ((<$>))
--   >>>
--   >>> let sort_snd (x,y) = (x, sort <$> y)
--   >>> rs <- makeResolvSeed defaultResolvConf
--   >>> let domain = BS.pack "example.com."
--   >>> withResolver rs $ \resolver -> sort_snd <$> lookupNS' resolver domain
--   ("example.com.",Right ["a.iana-servers.net.","b.iana-servers.net."])
--
lookupNS' :: Resolver -> Domain -> IO LookupResult
lookupNS' resolver domain = do
  answer_result <- lookupNS resolver domain
  auth_result <- lookupNSAuth resolver domain
  fmap pair_em $ return $ combine answer_result auth_result
  where
    pair_em :: a -> (Domain, a)
    pair_em = (,) domain

    combine :: (Either DNSError [Domain])
            -> (Either DNSError [Domain])
            -> (Either DNSError [Domain])
    combine e1 e2 = do
      l1 <- e1
      l2 <- e2
      return (l1 ++ l2)
