1 files changed, 45 insertions, 15 deletions

diff --git a/cokiki.hs b/cokiki.hs
index 2cb6491..daf2be5 100644
--- a/cokiki.hs
+++ b/cokiki.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE CPP #-}
 import Control.Applicative
+import Control.DeepSeq
 import Data.Bool
 import Data.Char
 import Data.Maybe
@@ -10,6 +12,7 @@ import System.Directory
 import System.Environment
 import System.IO
 import System.Posix.User
+import CommandLine
 
 usage = unlines
     [ "cokiki <command>"
@@ -20,37 +23,62 @@ usage = unlines
     , "  strongswan"
     ]
 
+ㄧchroot = pure (\r a -> slash r a) <*> arg "--chroot" <|> pure id
+ where
+    slash :: String -> String -> String
+    slash "/"    ('/':xs) = '/':xs
+    slash ""     ('/':xs) = '/':xs
+    slash ""     xs       = '/':xs
+    slash (y:ys) xs       = y:slash ys xs
+
 main = do
-    args <- getArgs
+    (cmd,args) <- splitAt 1 <$> getArgs
     uid <- getEffectiveUserID
     let whenRoot action
             | uid==0    = action
-            | otherwise = hPutStrLn stderr "operaiton requires root."
-    case args of
-        ["ssh-client"] -> sshClient uid id
-        ["ssh-server"] -> whenRoot sshServer
-        ["strongswan"] -> whenRoot strongswan
-        _   -> hPutStr stderr usage
+            | otherwise = hPutStrLn stderr "operation requires root."
+    let sel = case cmd of
+                ["ssh-client"] -> pure (sshClient uid) <*> ㄧchroot
+                ["ssh-server"] -> pure (whenRoot sshServer)
+                ["strongswan"] -> pure (whenRoot strongswan)
+                _   -> pure $ hPutStr stderr usage
+        spec = fancy [("--chroot",1)] [] ""
+    case runArgs (parseInvocation spec args) sel of
+        Left e   -> hPutStrLn stderr $ usageErrorMessage e
+        Right io -> io
 
 maybeReadFile :: FilePath -> IO (Maybe L.ByteString)
 maybeReadFile path = do
     doesFileExist path >>= bool (return Nothing) (Just <$> L.readFile path)
 
-sshClient 0 root = do
+sshClient uid root = do
     -- /etc/ssh/ssh_config    <--  'GlobalKnownHostsFile /var/cache/kiki/ssh_known_hosts'
     sshconfig <- parseSshConfig . fromMaybe "" <$> maybeReadFile (root "/etc/ssh/ssh_config")
     let (ps,qs) = sshSplitAtDirective "GlobalKnownHostsFile" sshconfig
+        myWriteFile f bs = do
+            hPutStrLn stderr $ "Writing "++f
+            -- L8.putStr bs
+            L8.writeFile f bs
     sshconfig' <-
         case qs of
-            d:ds -> error "modify GlobalKnownHostsFile"
-            []   -> error "append GlobalKnownHostsFile"
-    L.writeFile (root "/var/cache/kiki/ssh_known_hosts") $ unparseSshConfig sshconfig'
+            d:ds | elem "/var/cache/kiki/ssh_known_hosts" d
+                 -> do hPutStrLn stderr "ssh-client already configured."
+                       return Nothing
+            d:ds -> do hPutStrLn stderr "modifying GlobalKnownHostsFile not implemented"
+                       return Nothing
+            []   -> do
+                       let stmt = L8.unwords ["GlobalKnownHostsFile"
+                                             ,  "/var/cache/kiki/ssh_known_hosts"
+                                             ,  "/etc/ssh/ssh_known_hosts"
+                                             ,  "/etc/ssh/ssh_known_hosts2"
+                                             ]
+                       return $ Just (sshconfig ++ parseSshConfig stmt)
+    -- sshconfig' `deepseq` return () -- force lazy input
+    maybe (return ()) (myWriteFile (root "/etc/ssh/ssh_config") . unparseSshConfig) sshconfig'
     -- /root/.gnupg/...       <--  contains known hosts from /etc/ssh/ssh_known_hosts
-    Kiki.refresh
-    -- /var/cache/kiki/ssh_known_hosts        <--  contains known hosts from /root/.gnupg/...
-
 
-sshClient uid root = return ()
+    -- Kiki.refresh
+    -- /var/cache/kiki/ssh_known_hosts        <--  contains known hosts from /root/.gnupg/...
 
 sshServer = do
     -- /etc/ssh/sshd_config <-- 'HostKey /var/cache/kiki/ssh_host_ecdsa_key' etc.
@@ -84,6 +112,8 @@ sshIsDirective d ls =
     isSpaceTok "" = True
     isSpaceTok b  = isSpace $ L8.head b
 
+#if !MIN_VERSION_base(4,8,0)
 bool :: a -> a -> Bool -> a
 bool f _ False = f
 bool _ t True  = t
+#endif