1 files changed, 76 insertions, 17 deletions

diff --git a/Hosts.hs b/Hosts.hs
index 669fd09..e53d441 100644
--- a/Hosts.hs
+++ b/Hosts.hs
@@ -12,11 +12,14 @@ module Hosts
     , encode
     , decode
     , diff
+    , plus
+    , filterAddrs
+    , namesForAddress
     ) where
 
 import Data.Maybe
 import Data.Monoid ( (<>) )
-import Data.List (foldl')
+import Data.List as List (foldl', (\\) )
 import Data.Ord
 import Data.Char (isSpace)
 import qualified Data.Map as Map
@@ -77,9 +80,9 @@ empty = Hosts { lineCount = 0
 
 parseHosts fname = do
     input <- L.readFile fname
-    decode input
+    return $ decode input
 
-decode input = do
+decode input = 
     let ls = L.lines input
         ans = map (\l->(parseLine l,l)) ls
         hosts = foldl' upd empty ans
@@ -96,7 +99,7 @@ decode input = do
                 maybeInsert m x = maybe m
                                         (\x->Map.insert x count m)
                                         x
-    return hosts
+    in hosts
 
 
 hasName :: L.ByteString -> Hosts -> Bool
@@ -112,8 +115,20 @@ scrubName f line = line'
                         where oo = (.) . (.)
     (a,ws') = splitAt 2 ws
     ws'' = f ws'
-    line' = if null ws'' then "# " <> line
-                         else L.concat (a ++ ws'') <> ign
+    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
 
@@ -128,23 +143,13 @@ assignName0 iscannon addr name hosts = hosts'
             let hosts0 = -- remove name if it's present
                          maybe hosts (removeName hosts) ns
                 hosts1 = -- insert name, or add new line
-                         maybe (newLine hosts0) (appendName hosts0) a
+                         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
-    appendName 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
     newLine hosts = hosts
         { lineCount = cnt
         , numline = Map.insert cnt line $ numline hosts
@@ -159,6 +164,17 @@ 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
@@ -175,3 +191,46 @@ diff as bs = cs
         [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