1 files changed, 163 insertions, 0 deletions

diff --git a/Hosts.hs b/Hosts.hs
new file mode 100644
index 0000000..4f8788f
--- /dev/null
+++ b/Hosts.hs
@@ -0,0 +1,163 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE OverloadedStrings   #-}
+module Hosts
+    ( Hosts
+    , assignName
+    , assignNewName
+    , inet_pton
+    , inet_ntop 
+    , empty
+    , hasName
+    , hasAddr
+    ) where
+
+import Data.Maybe
+import Data.Monoid ( (<>) )
+import Data.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 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 hosts = L.unpack . L.unlines . map snd . Map.assocs $ numline hosts
+
+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
+    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
+    return 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 "# " <> line
+                         else L.concat (a ++ ws'') <> ign
+
+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 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
+        , 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
+
+{-
+main = do
+    args <- getArgs
+    let fname = args !! 0
+    p <- parseHosts fname
+    let addr = (fromJust $ inet_pton "fdee:0abe:1f80:31c7:d1af:bce0:0f6c:91d2")
+        p'   = assignName addr "bigshift" p
+        p''  = assignNewName addr "poopy" p'
+    putStr $ show p''
+    return ()
+-}