path: root/lib/Hosts.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings   #-}
module Hosts
    ( Hosts
    , assignName
    , assignName'
    , assignNewName
    , removeName
    , 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.Applicative ( (<$>), (<*>) )
import Control.Monad (mplus)
import Network.Socket
import ControlMaybe ( handleIO_ )

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 :: Hosts -> L.ByteString
encode = L.unlines . map snd . Map.assocs . numline

parseLine :: L.ByteString -> (Maybe SockAddr, [L.ByteString])
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
empty = Hosts { lineCount = 0
              , numline = Map.empty
              , addrnum = Map.empty
              , namenum = Map.empty
              }

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

decode :: L.ByteString -> Hosts
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 ::
  ([L.ByteString] -> [L.ByteString]) -> L.ByteString -> L.ByteString
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 :: SockAddr -> L.ByteString -> Hosts -> Hosts
assignName addr name hosts = assignName' False addr name hosts

chaddr :: Int -> SockAddr -> Hosts -> Hosts
chaddr n addr hosts =
    hosts { addrnum = Map.insert addr n (addrnum hosts)
          , numline = Map.adjust re n (numline hosts) }
 where
    re line = if length a==2
                then L.pack (inet_ntop addr) <> " " <> L.concat ws' <> ign
                else line
        where (x,ign) = L.break (=='#') line
              ws  = L.groupBy ( (==EQ) `oo` comparing isSpace) x
                                  where oo = (.) . (.)
              (a,ws') = splitAt 2 ws

isLonerName :: L.ByteString -> Bool
isLonerName line = length ws' <= 2
        where (x,_) = L.break (=='#') line
              ws  = L.groupBy ( (==EQ) `oo` comparing isSpace) x
                                  where oo = (.) . (.)
              (_,ws') = splitAt 2 ws

scrubTrailingEmpties :: Hosts -> Hosts
scrubTrailingEmpties hosts =
    hosts { lineCount = cnt'
          , numline = foldl' (flip Map.delete) (numline hosts) es
          }
 where
    cnt = lineCount hosts
    es = takeWhile (\n -> Map.lookup n (numline hosts) == Just "")
            $ [cnt,cnt-1..]
    cnt' = cnt - length es

cannonizeName :: L.ByteString -> L.ByteString -> L.ByteString
cannonizeName name line = scrubName f line
 where
    f ws = [name," "] ++ pre ++ drop 2 rst
      where
        (pre,rst) = break (==name) ws

removeName :: L.ByteString -> Hosts -> Hosts
removeName name hosts = hosts'
 where
    hosts' = scrubTrailingEmpties (maybe hosts (removeName0 name hosts) ns)
    ns = Map.lookup name (namenum hosts) 

        
removeName0 :: Foldable t =>
               L.ByteString -> Hosts -> t Int -> Hosts
removeName0 name 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

assignName' :: Bool -> SockAddr -> L.ByteString -> Hosts -> Hosts
assignName' iscannon addr name hosts = hosts'
 where
    ns = Map.lookup name (namenum hosts) 
    a = Map.lookup addr (addrnum hosts)
    canonize numline n = Map.adjust (cannonizeName name) n numline
    hosts' = do
        if (== Just True) $ elem <$> a <*> ns
         then if not iscannon then hosts -- address already has name, nothing to do
                              else hosts { numline = foldl' canonize (numline hosts) $ fromJust ns}
         else 
            let hosts0 = -- remove name if it's present
                         scrubTrailingEmpties $ maybe hosts (removeName0 name hosts) ns
                ns' = fmap (filter $
                                isLonerName
                                . fromJust
                                . (\n -> Map.lookup n (numline hosts)))
                            ns
                        >>= listToMaybe
                hosts1 = -- insert name, or add new line
                         maybe (maybe (newLine hosts0)
                                      (\n -> chaddr n addr $ appendName iscannon name hosts0 n)
                                      ns')
                               (appendName iscannon name hosts0)
                               a
            in hosts1
    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 :: SockAddr -> L.ByteString -> Hosts -> Hosts
assignNewName addr name hosts = 
    if hasName name hosts then hosts
                          else assignName' True addr name hosts

appendName :: Bool -> L.ByteString -> Hosts -> Int -> 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 :: SockAddr -> Hosts -> (Int, [L.ByteString])
_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 :: Hosts -> Hosts -> Hosts
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