module Network.DNS.RBL ( lookup_single )
where

import qualified Data.ByteString.Char8 as BS (
  append,
  intercalate,
  pack,
  split )
import Data.IP ( IPv4 )
import Data.List ( intercalate )
import Data.Maybe ( catMaybes )
import Network.DNS (
  DNSError,
  Resolver,
  defaultResolvConf,
  lookupA,
  makeResolvSeed,
  withResolver )

import Network.DNS.RBL.Domain ( UserDomain )
import Network.DNS.RBL.IPv4Pattern ( addresses )
import Network.DNS.RBL.Pretty ( Pretty(..) )
import Network.DNS.RBL.Site ( Site(..), Weight(..) )

-- | A dumb wrapper around an IP address input as a 'String'.
--
newtype Host = Host String
instance Pretty Host where pretty_show (Host h) = h


-- | In most applications, you will want to know /why/ a host is
--   blacklisted. This type stores those details: the site on which
--   the host was listed, and the return codes that we got back from
--   the blacklist.
data ListingDetails = ListingDetails Site [IPv4]


-- | Create a nice \"error\" message from a host and the details of
--   why it was listed.
--
listing_message :: Host -> ListingDetails -> String
listing_message host (ListingDetails (Site d _ w) codes) =
  "host " ++ (pretty_show host) ++ " " ++
  "listed on " ++ (pretty_show d) ++ " " ++
  "with return code(s) " ++ return_codes ++ " " ++
  "having weight " ++ (pretty_show w)
  where
    return_codes :: String
    return_codes = intercalate "," (map show codes)


-- | This code is stolen from 'Network.DNS.lookupRDNS', which I'm
--   pretty sure I wrote anyway.
--
dnslookup :: Resolver -> UserDomain -> Host -> IO (Either DNSError [IPv4])
dnslookup rlv rbl (Host ip) = lookupA rlv dom
  where
    -- ByteString constants.
    dot = BS.pack "."
    suffix = BS.pack $ "." ++ (pretty_show rbl)

    octets = BS.split '.' (BS.pack ip)
    reverse_ip = BS.intercalate dot (reverse octets)
    dom = reverse_ip `BS.append` suffix



-- | Look up the given @host@ on a single RBL, and return the number
--   of \"points\" it scores. If the host is listed, it should score the
--   weight of the given blacklist. Otherwise, it scores zero.
--
lookup_single :: Host -> Site -> IO Int
lookup_single host site = do
  rs <- makeResolvSeed defaultResolvConf
  withResolver rs $ \resolver -> lookup_single' resolver host site


-- | See 'lookup_single'.
--
--   The \"prime\" version of this function takes a pre-existing
--   resolver as an argument for performance reasons.
--
lookup_single' :: Resolver -> Host -> Site -> IO Int
lookup_single' resolver host (Site d p (Weight w)) = do
  response <- dnslookup resolver d host
  case response of
    Left _ -> return 0 -- Not listed
    Right ipv4s ->
      case p of
        -- No pattern given, but we got a hit.
        Nothing -> return w
        Just pat -> do
          let ipv4_strings = map show ipv4s
          let codes = addresses pat
          let hits = map (`elem` codes) ipv4_strings
          if or hits -- if any of the returned addresses match the pattern
          then return w
          else return 0


-- | Look up the given @host@ on all of the white/blacklists contained
--   in @sites@. Return the total \"score\" for the @host@; that is,
--   the sum of the weights for each site whose pattern the host
--   matches.
--
lookup_simple :: Host -> [Site] -> IO Int
lookup_simple host sites = do
  rs <- makeResolvSeed defaultResolvConf
  withResolver rs $ \resolver -> do
    results <- mapM (lookup_single' resolver host) sites
    return $ sum results

-- | Look up the given @host@ on all of the white/blacklists contained
--   in @sites@. If the results, multiplied by their weights, add up to
--   (at least) @threshold@, then @True@ is returned; that is, the
--   @host@ is \"listed\". Otherwise, @False@ is returned.
--
lookup_threshold :: Host -> [Site] -> Int -> IO Bool
lookup_threshold host sites threshold = do
  score <- lookup_simple host sites
  if score >= threshold
  then return True
  else return False


-- | See 'lookup_single_details'. The \"prime\" version here takes an
--   additional @resolver@ argument for performance reasons.
--
lookup_single_details' :: Resolver
                       -> Host
                       -> Site
                       -> IO (Maybe ListingDetails)
lookup_single_details' resolver host site@(Site d p (Weight w)) = do
  response <- dnslookup resolver d host
  case response of
    Left _ -> return Nothing -- Not listed, no error to report
    Right ipv4s ->
      case p of
        -- No pattern given, but we got a hit.
        Nothing -> return $ Just (ListingDetails site ipv4s)
        Just pat -> do
          let ipv4_strings = map show ipv4s
          let codes = addresses pat
          let hits = map (`elem` codes) ipv4_strings
          if or hits -- if any of the returned addresses match the pattern
          then return $ Just (ListingDetails site ipv4s)
          else return Nothing


-- | Look up a single...
lookup_single_details :: Host
                      -> Site
                      -> IO (Maybe ListingDetails)
lookup_single_details host site = do
  rs <- makeResolvSeed defaultResolvConf
  withResolver rs $ \resolver -> lookup_single_details' resolver host site


lookup_details :: Host -> [Site] -> IO [ListingDetails]
lookup_details host sites = do
  rs <- makeResolvSeed defaultResolvConf
  withResolver rs $ \resolver -> do
    results <- mapM (lookup_single_details' resolver host) sites
    return $ catMaybes results