1 files changed, 25 insertions, 3 deletions

diff --git a/KnownHosts.hs b/KnownHosts.hs
index e04ad90..e19e68e 100644
--- a/KnownHosts.hs
+++ b/KnownHosts.hs
@@ -1,9 +1,17 @@
 {-# LANGUAGE OverloadedStrings #-}
-module KnownHosts where
+module KnownHosts
+    ( KnownHosts(..)
+    , decode
+    , encode
+    , assignKey
+    , parseLine
+    , serializeLine ) where
 
 import qualified Data.ByteString.Lazy.Char8 as L
 import Data.Char
 import Data.List
+import Data.Ord
+import Data.Maybe
 import qualified SSHKey as SSH
 
 data KnownHostsEntry = KnownHostsEntry
@@ -61,14 +69,28 @@ data KnownHosts = KnownHosts
     , khentries :: [KnownHostsEntry]
     }
 
+renum kh = kh { khentries = zipWith setnum [1..] (khentries kh) }
+        where setnum n e = e { eLineno=n }
+
 decode :: L.ByteString -> KnownHosts
 decode input = renum $ foldr grok (KnownHosts 0 []) ls
  where
     ls = L.lines input
     grok x (KnownHosts cnt es) =
         KnownHosts (cnt+1) (parseLine 0 x:es)
-    renum kh = kh { khentries = zipWith setnum [1..] (khentries kh) }
-        where setnum n e = e { eLineno=n }
 
 encode :: KnownHosts -> L.ByteString
 encode kh = L.unlines $ map serializeLine (khentries kh)
+
+assignKey :: SSH.Key -> L.ByteString -> KnownHosts -> KnownHosts
+assignKey key name kh = renum KnownHosts { khcount = khcount kh - length ds
+                                         , khentries = es }
+ where
+    (ms,ns) = partition (\e -> eKey e == Just key) $ khentries kh
+    ms' = map (\e -> e { eHosts = eHosts e `union` [name]}) ms
+    ns' = map (\e -> e { eHosts = eHosts e \\ if isJust (eKey e) then [name]
+                                                                 else []
+                       })
+              ns
+    (ds,ns'') = partition (\e -> null (eHosts e)) ns'
+    es = sortBy (comparing eLineno) $ ms' ++ ns''