path: root/Hosts.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings   #-}
module Hosts
    ( Hosts
    , assignName
    , assignNewName
    , inet_pton
    , inet_ntop 
    , empty
    , hasName
    , hasAddr
    , encode
    , decode
    , diff
    , plus
    , filterAddrs
    , namesForAddress
    ) where

import Data.Maybe
import Data.Monoid ( (<>) )
import Data.List as List (foldl', (\\) )
import Data.Ord
import Data.Char (isSpace)
import qualified Data.Map as Map
import           Data.Map (Map)
import qualified Data.ByteString.Lazy.Char8 as L
import System.IO.Unsafe (unsafePerformIO)
import Control.Exception as Exception (IOException(..),catch)
import Control.Applicative ( (<$>), (<*>) )
import Control.Monad (mplus)
import Network.Socket

handleIO_ h a = Exception.catch a (\(_ :: IOException) -> h)

inet_pton :: String -> Maybe SockAddr
inet_pton p = n
 where
    n = unsafePerformIO $ do
            handleIO_ (return Nothing) $ do
            info <- getAddrInfo safe_hints (Just p) Nothing
            return $ fmap addrAddress $ listToMaybe info
    safe_hints = Just $ defaultHints { addrFlags=[AI_NUMERICHOST] }

inet_ntop :: SockAddr -> String
inet_ntop n = p
 where
    p = case show n of
            '[':xs ->  fst $ break (==']') xs
            xs     ->  fst $ break (==':') xs


data Hosts = Hosts
    { lineCount :: Int
    , numline :: Map Int L.ByteString
    , namenum :: Map L.ByteString [Int]
    , addrnum :: Map SockAddr Int
    }

instance Show Hosts where
    show = L.unpack . encode

encode = L.unlines . map snd . Map.assocs . numline

parseLine s = (addr,names)
 where
    (addr0,names) = splitAt 1 $ L.words (uncom s)
    addr = do
        a <- fmap L.unpack $ listToMaybe addr0
        n <- inet_pton a
        return $ n -- inet_ntop n
        
    uncom s = fst $ L.break (=='#') s

empty = Hosts { lineCount = 0
              , numline = Map.empty
              , addrnum = Map.empty
              , namenum = Map.empty
              } 

parseHosts fname = do
    input <- L.readFile fname
    return $ decode input

decode input = 
    let ls = L.lines input
        ans = map (\l->(parseLine l,l)) ls
        hosts = foldl' upd empty ans
        upd hosts ((addr,names),line) = hosts
            { lineCount = count
            , numline = Map.insert count line (numline hosts)
            , addrnum = maybeInsert (addrnum hosts) addr
            , namenum = foldl' (\m x->Map.alter (cons count) x m)
                               (namenum hosts)
                               names
            }
          where count = lineCount hosts + 1
                cons v xs = Just $ maybe [v] (v:) xs
                maybeInsert m x = maybe m
                                        (\x->Map.insert x count m)
                                        x
    in hosts


hasName :: L.ByteString -> Hosts -> Bool
hasName name hosts = Map.member name $ namenum hosts

hasAddr :: SockAddr -> Hosts -> Bool
hasAddr addr hosts = Map.member addr $ addrnum hosts

scrubName f line = line'
 where
    (x,ign) = L.break (=='#') line
    ws  = L.groupBy ( (==EQ) `oo` comparing isSpace) x
                        where oo = (.) . (.)
    (a,ws') = splitAt 2 ws
    ws'' = f ws'
    line' = if null ws'' 
                then if length a==2 then "# " <> L.concat a <> ign
                                    else line
                else if length a==2 
                    then L.concat (a ++ ws'') <> ign
                    else let vs  = L.groupBy ( (==EQ) `oo` comparing isSpace) $ L.dropWhile isSpace
                                                                              $ L.tail ign
                                                 where oo = (.) . (.)
                             (a',vs') = splitAt 2 vs
                             vs'' = L.concat vs'
                             vs''' = if L.take 1 vs'' `elem` ["#",""] 
                                        then vs''
                                        else "# " <> vs''
                         in L.concat (a'++ws'') <> vs'''

assignName addr name hosts = assignName0 False addr name hosts

assignName0 iscannon addr name hosts = hosts'
 where
    ns = Map.lookup name (namenum hosts) 
    a = Map.lookup addr (addrnum hosts)
    hosts' = do
        if (== Just True) $ elem <$> a <*> ns
         then hosts -- address already has name, nothing to do
         else 
            let hosts0 = -- remove name if it's present
                         maybe hosts (removeName hosts) ns
                hosts1 = -- insert name, or add new line
                         maybe (newLine hosts0) (appendName iscannon name hosts0) a
            in hosts1
    removeName hosts nums = hosts 
        { namenum = Map.delete name (namenum hosts)
        , numline = foldl' scrub (numline hosts) nums
        }
        where scrub m num = Map.adjust (scrubName $ filter (/=name)) num m
    newLine hosts = hosts
        { lineCount = cnt
        , numline = Map.insert cnt line $ numline hosts
        , addrnum = Map.insert addr cnt $ addrnum hosts
        , namenum = Map.alter (cons cnt) name $ namenum hosts
        }
        where cnt = lineCount hosts + 1
              line = L.pack (inet_ntop addr) <> " " <> name
              cons v xs = Just $ maybe [v] (v:) xs

assignNewName addr name hosts = 
    if hasName name hosts then hosts
                          else assignName0 True addr name hosts

appendName iscannon name hosts num = hosts
        { numline = Map.adjust (scrubName f) num (numline hosts)
        , namenum = Map.alter (cons num) name (namenum hosts)
        }
        where f ws = if iscannon 
                      then [name, " "] ++ ws
                      else let rs       = reverse ws
                               (sp,rs') = span (L.any isSpace) rs
                           in reverse $ sp ++ [name," "] ++ rs'
              cons v xs = Just $ maybe [v] (v:) xs

-- Returns a list of bytestrings intended to show the
-- differences between the two host databases.  It is
-- assumed that no lines are deleted, only altered or
-- appended.
diff :: Hosts -> Hosts -> [L.ByteString]
diff as bs = cs
 where
    [as',bs'] = map (L.lines . Hosts.encode) [as,bs]
    ext xs = map Just xs ++ repeat Nothing
    ds = takeWhile (isJust . uncurry mplus) $ zip (ext as') (ext bs')
    es = filter (uncurry (/=)) ds
    cs = do
        (a,b) <- es
        [a,b] <- return $ map maybeToList [a,b]
        fmap ("- " <>) a ++ fmap ("+ " <>) b

namesForAddress :: SockAddr -> Hosts -> [L.ByteString]
namesForAddress addr hosts = snd $ _namesForAddress addr hosts

_namesForAddress addr (Hosts {numline=numline, addrnum=addrnum}) = ns
 where
    ns = maybe (-1,[]) id $ do
        n <- Map.lookup addr addrnum
        line <- Map.lookup n numline
        return (n, snd $ parseLine line)


plus a b = Map.foldlWithKey' mergeAddr a (addrnum b)
 where
    mergeAddr a addr bnum = a'
     where
        (anum,ns) = _namesForAddress addr a
        bs = maybe [] (List.\\ ns) $ do
                 line <- Map.lookup bnum (numline b)
                 return . snd $ parseLine line
        a' = if anum/=(-1) then foldl' app a $ reverse bs
                           else newLine a
        app a b = appendName True b a anum -- True to allow b to reassign cannonical name
        newLine hosts = hosts
            { lineCount = cnt
            , numline = Map.insert cnt line $ numline hosts
            , addrnum = Map.insert addr cnt $ addrnum hosts
            , namenum = foldl' updnamenum (namenum hosts) bs
            }
            where cnt = lineCount hosts + 1
                  line = L.pack (inet_ntop addr) <> " " <> L.intercalate " " bs
                  cons v xs = Just $ maybe [v] (v:) xs
                  updnamenum m name = Map.alter (cons cnt) name m

filterAddrs :: (SockAddr -> Bool) -> Hosts -> Hosts
filterAddrs pred hosts = hosts'
 where
    als = Map.toList (addrnum hosts)
    nl  = foldl' f (numline hosts) als
    f m (addr,num) = if pred addr
                        then m
                        else Map.adjust (scrubName $ const []) num m
    lines = L.unlines . Map.elems $ nl
    hosts' = decode lines