{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
module Kiki where

import Control.Applicative
import Control.Arrow
import Control.Monad
import Data.ASN1.BinaryEncoding
import Data.ASN1.Encoding
import Data.ASN1.Types
import Data.Binary
import Data.List
import Data.Maybe
import Data.Monoid
import Data.OpenPGP
import Data.OpenPGP.Util
import Data.Ord
import System.Directory
import System.FilePath.Posix
import System.IO
import System.Posix.User
import System.Process
import qualified Codec.Binary.Base64 as Base64
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Char8 as Char8
import qualified Data.Map.Strict as Map
import qualified SSHKey as SSH

import CommandLine
import KeyRing

-- |
-- Regenerate /var/cache/kiki
refresh :: (FilePath -> FilePath) -> CommonArgsParsed -> IO ()
refresh root homepass = do
    let homepass' = homepass { cap_homespec = fmap root (cap_homespec homepass) }
    KikiResult r report <- runKeyRing $ minimalOp homepass'
    let mroot = case root "" of
                    "/" -> Nothing
                    ""  -> Nothing
                    pth -> Just pth
    case r of
        KikiSuccess rt -> refreshCache rt mroot
        _ -> return () -- XXX: silent fail?

data CommonArgsParsed = CommonArgsParsed { cap_homespec :: Maybe String, cap_passfd :: Maybe InputFile }

streaminfo :: StreamInfo
streaminfo = StreamInfo
    { fill = KF_None
    , spill = KF_None
    , typ = KeyRingFile
    , initializer = NoCreate
    , access = AutoAccess
    , transforms = []
    }

minimalOp :: CommonArgsParsed -> KeyRingOperation
minimalOp cap = op
 where
    streaminfo = StreamInfo { fill = KF_None
                            , typ = KeyRingFile
                            , spill = KF_All
                            , initializer = NoCreate
                            , access = AutoAccess
                            , transforms = []
                            }
    op = KeyRingOperation
        { opFiles = Map.fromList $
            [ ( HomeSec, streaminfo { access = Sec })
            , ( HomePub, streaminfo { access = Pub })
            ]
        , opPassphrases = do pfile <- maybeToList (cap_passfd cap)
                             return $ PassphraseSpec Nothing Nothing pfile
        , opTransforms = []
        , opHome = cap_homespec cap
        }

run :: [String] -> Args (IO ()) -> IO ()
run args x =
    case runArgs (parseInvocation (uncurry fancy kikiOptions "") args) x of
        Left e   -> hPutStrLn stderr $ usageErrorMessage e
        Right io -> io

importAndRefresh :: (FilePath -> FilePath) -> CommonArgsParsed -> IO ()
importAndRefresh root cmn = do
    let rootdir = do guard (root "x" /= "x")
                     Just $ root ""

    me <- getEffectiveUserID

    let noChrootArg = rootdir == Nothing
        bUnprivileged = (me/=0) && noChrootArg
    if rootdir==Just "" then error "--chroot requires an argument" else do

    let homespec = mplus (slash <$> rootdir <*> cap_homespec cmn)
                         (fmap (++"/root/.gnupg") rootdir)
        sshkeygen size = Just $ concat [ "mkdir -p \"$(dirname $file)\" && "
                                       , "ssh-keygen -P \"\" -q -f $file -b "
                                       , show size ]
        mkdirFor path = do
            let dir = takeDirectory path
            -- putStrLn $ "mkdirFor " ++ show dir
            createDirectoryIfMissing True dir
        -- ssl = Just "mkdir -p \"$(dirname $file)\" && openssl genrsa -out $file 1024"
    (home,secring,pubring,mbwk) <- unconditionally $ getHomeDir homespec
    osHomeDir <- if bUnprivileged then getHomeDirectory else return "/root"

    -- Generate secring.gpg if it does not exist...
    gotsec <- doesFileExist secring
    when (not gotsec) $ do
        {-  ssh-keygen to create master key...
        let mkpath = home ++ "/master-key"
        mkdirFor mkpath
        e <- systemEnv [ ("file",mkpath) ] (fromJust $ sshkeygen 4096)
        case e of
            ExitFailure num -> error "ssh-keygen failed to create master key"
            ExitSuccess     -> return ()
        [PEMPacket mk] <- readSecretPEMFile (ArgFile mkpath)
        writeInputFileL (InputFileContext secring pubring)
                        HomeSec
                        ( encode $ Message [mk { is_subkey = False }] )
        -}
        master <- (\k -> k { is_subkey = False }) <$> generateKey (GenRSA $ 4096 `div` 8 )
        mkdirFor secring
        writeInputFileL (InputFileContext secring pubring)
                        HomeSec
                        $ encode $ Message [master { is_subkey = False}]

    gotpub <- doesFileExist pubring
    when (not gotpub) $ do
        mkdirFor pubring
        writeInputFileL (InputFileContext secring pubring)
                        HomePub
                        ( encode $ Message [] )

    -- Old paths..
    --
    -- Private
    -- pem tor        /var/lib/tor/samizdat/private_key
    -- pem ssh-client %(home)/.ssh/id_rsa
    -- pem ssh-server /etc/ssh/ssh_host_rsa_key
    -- pem ipsec      /etc/ipsec.d/private/%(onion).pem

    -- Public
    -- ssh-client %(home)/.ssh/id_rsa.pub
    -- ssh-server /etc/ssh/ssh_host_rsa_key.pub
    -- ipsec      /etc/ipsec.d/certs/%(onion).pem

    -- First, we ensure that the tor key exists and is imported
    -- so that we know where to put the strongswan key.
    let passfd = cap_passfd cmn
        strm = StreamInfo { typ = KeyRingFile
                          , fill = KF_None
                          , spill = KF_All
                          , access = AutoAccess
                          , initializer = NoCreate
                          , transforms = [] }
        buildStreamInfo rtyp ftyp = StreamInfo { typ = ftyp
                                               , fill = rtyp
                                               , spill = KF_All
                                               , access = AutoAccess
                                               , initializer = NoCreate
                                               , transforms = [] }
        peminfo bits usage =
                  StreamInfo { typ = PEMFile
                             , fill = KF_None -- KF_Match usage
                             , spill = KF_Match usage
                             , access = Sec
                             , initializer = Internal (GenRSA $ bits `div` 8)
                             , transforms = []
                             }
        sshcpath = fromMaybe "" rootdir ++ osHomeDir ++ ".ssh/id_rsa"
        sshspath = fromMaybe "" rootdir ++ "/etc/ssh/ssh_host_rsa_key"
        op = KeyRingOperation
             { opFiles = Map.fromList $
                [ ( HomeSec, buildStreamInfo KF_All KeyRingFile )
                , ( HomePub, (buildStreamInfo KF_All KeyRingFile) { access = Pub } )
                , ( Generate 0 (GenRSA (1024 `div` 8)), strm { spill = KF_Match "tor" })
                , ( Generate 1 (GenRSA (1024 `div` 8)), strm { spill = KF_Match "ipsec" })
                , ( ArgFile sshcpath, (peminfo 2048 "ssh-client") )
                , ( ArgFile sshspath, (peminfo 2048 "ssh-server") )
                ]
             , opPassphrases = do pfd <- maybeToList passfd
                                  return $ PassphraseSpec Nothing Nothing pfd
             , opHome = homespec
             , opTransforms = []
             }
        -- doNothing = return ()
        nop = KeyRingOperation
            { opFiles = Map.empty
            , opPassphrases = do pfd <- maybeToList passfd
                                 return $ PassphraseSpec Nothing Nothing pfd
            , opHome=homespec, opTransforms = []
            }
    -- if bUnprivileged then doNothing else mkdirFor torpath
    KikiResult rt report <- runKeyRing (if bUnprivileged then nop else op)
    forM_ report $ \(fname,act) -> do
        putStrLn $ fname ++ ": " ++ reportString act
    rt <- case rt of
        BadPassphrase ->
            error "Operation requires correct passphrase. (Hint: Use --passphrase-fd=0 to input it on stdin.)"
        _ -> unconditionally $ return rt

    when (not bUnprivileged) $ refreshCache rt rootdir


refreshCache :: KeyRingRuntime -> Maybe FilePath -> IO ()
refreshCache rt rootdir = do

    let mkpath pth = fromMaybe "" rootdir ++ "/var/cache/kiki/"++pth

        write' wr f bs = do
            createDirectoryIfMissing True $ takeDirectory f
            wr f bs
        write = write' writeFile
        writeL = write' L.writeFile

    let names = do wk <- rtWorkingKey rt
                   -- XXX unnecessary signature check
                   return $ getHostnames (rtKeyDB rt Map.! keykey wk)
        bUnprivileged = False -- TODO
        oname = Char8.concat $ do
            (_,(os,_)) <- maybeToList names
            take 1 os
    fromMaybe (error "No working key.") $ do
    (wkaddr,_) <- names
    Just $ do
    if (oname == "") && (not bUnprivileged) then error "Missing tor key" else do
    -- sshcpathpub0 = fromMaybe "" rootdir ++ osHomeDir </> ".ssh" </> "id_rsa.pub"
    -- sshspathpub0 = fromMaybe "" rootdir ++ "/etc/ssh/ssh_host_rsa_key.pub"
    -- contactipsec0 = fromMaybe "" rootdir ++ "/etc/ipsec.d/certs/%(onion).pem"

    callCommand ("rm -rf "++ mkpath "*") -- clean up, in case gpg altered the keyring.

    -- Finally, export public keys if they do not exist.
    flip (maybe $ warn "missing working key?") (rtGrip rt) $ \grip -> do
    either warn (write $ mkpath "root/.ssh/id_rsa.pub")
            $ show_ssh' "ssh-client" grip (rtKeyDB rt)
    either warn (write $ mkpath "ssh_host_rsa_key.pub")
        $ show_ssh' "ssh-server" grip (rtKeyDB rt)
    either warn (write $ mkpath "ipsec.d/certs/" ++ Char8.unpack oname++".pem")
        $ show_pem' "ipsec" grip (rtKeyDB rt) pemFromPacket

    let cs = filter notme (Map.elems $ rtKeyDB rt)
        kk = keykey (fromJust $ rtWorkingKey rt)
        notme kd = keykey (keyPacket kd) /= kk

        installConctact :: KeyData -> IO Char8.ByteString
        installConctact kd = do
            -- The getHostnames command requires a valid cross-signed tor key
            -- for each onion name returned in (_,(ns,_)).
            let (addr,(ns,_)) = getHostnames kd
                contactname = fmap Char8.unpack $ listToMaybe ns -- only first onion name.
            flip (maybe $ return Char8.empty) contactname $ \contactname -> do

            let cpath = interp (Map.singleton "onion" contactname) "ipsec.d/certs/%(onion).pem"
                their_master = packet $ keyMappedPacket kd
                -- We find all cross-certified ipsec keys for the given cross-certified onion name.
                ipsecs :: [Packet]
                ipsecs = sortOn (Down . timestamp)
                         $ getCrossSignedSubkeys their_master (keySubKeys kd) "ipsec"
            bss <- forM (take 1 ipsecs) $ \k -> do
                let warn' x = warn x >> return Char8.empty
                flip (either warn') (pemFromPacket k :: Either String String) $ \pem -> do
                    write (mkpath cpath) pem
                    return $ strongswanForContact addr contactname
            return $ Char8.concat bss

    cons <- mapM installConctact cs
    writeL (mkpath "ipsec.conf") . Char8.unlines
          $ [ "conn %default"
            , "  ikelifetime=60m"
            , "  keylife=20m"
            , "  rekeymargin=3m"
            , "  keyingtries=%forever"
            , "  keyexchange=ikev2"
            , "  dpddelay=10s"
            , "  dpdaction=restart"
            , "  left=%defaultroute"
            , "  leftsubnet=" <> Char8.pack (showA wkaddr) <> "/128"
            , "  leftauth=pubkey"
            , "  leftid=" <> Char8.pack (showA wkaddr)
            , "  leftrsasigkey=" <> oname
            , "  leftikeport=4500"
            , "  rightikeport=4500"
            , "  right=%any"
            , "  rightauth=pubkey"
            , "  type=tunnel"
            , "  auto=route"
            , ""
            ] ++ filter (not . Char8.null) cons
    return ()

strongswanForContact addr oname = Char8.unlines
    [ "conn " <> p oname
    , "  right=%" <> p oname <> ".ipv4"
    , "  rightsubnet=" <> p (showA addr) <> "/128"
    , "  rightauth=pubkey"
    , "  rightid=" <> p (showA addr)
    , "  rightrsasigkey=" <> p (oname) <> ".pem"
    ]
    where p = Char8.pack

-- conn hiotuxliwisbp6mi.onion
--   right=%hiotuxliwisbp6mi.onion.ipv4
--   rightsubnet=fdcc:76c8:cb34:74e6:2aa3:cb39:abc8:d403/128
--   rightauth=pubkey
--   rightid=fdcc:76c8:cb34:74e6:2aa3:cb39:abc8:d403
--   rightrsasigkey=hiotuxliwisbp6mi.onion.pem

showA addr = if null bracket then pre else drop 1 pre
    where (pre,bracket) = break (==']') (show addr)


#if !MIN_VERSION_base(4,8,0)
sortOn :: Ord b => (a -> b) -> [a] -> [a]
sortOn f =
    map snd . sortBy (comparing fst) . map (\x -> let y = f x in y `seq` (y, x))
#endif

pemFromPacket k = do
    let rsa = pkcs8 . fromJust $ rsaKeyFromPacket k
        der = encodeASN1 DER (toASN1 rsa [])
        qq  = Base64.encode (L.unpack der)
    return $
        writePEM "PUBLIC KEY" qq -- ("TODO "++show keyspec)

show_pem keyspec wkgrip db = either warn putStrLn $ show_pem' keyspec wkgrip db pemFromPacket

show_pem' keyspec wkgrip db keyfmt = do
    let s = parseSpec wkgrip keyspec
    flip (maybe . Left $ keyspec ++ ": not found")
         (selectPublicKey s db)
         keyfmt

warn str = hPutStrLn stderr str

show_ssh keyspec wkgrip db = either warn putStrLn $ show_ssh' keyspec wkgrip db

show_ssh' keyspec wkgrip db = do
    let s = parseSpec wkgrip keyspec
    flip (maybe . Left $ keyspec ++ ": not found")
         (selectPublicKey s db)
         $ return . sshblobFromPacket

-- |
-- interpolate %var patterns in a string.
interp vars raw = es >>= interp1
 where
    gs = groupBy (\_ c -> c/='%') raw
    es = dropWhile null $ gobbleEscapes ("":gs)
            where gobbleEscapes :: [String] -> [String]
                  gobbleEscapes (a:"%":b:bs) = (a++b) : gobbleEscapes bs
                  gobbleEscapes (g:gs)       = g : gobbleEscapes gs
                  gobbleEscapes []           = []
    interp1 ('%':'(':str) = fromMaybe "" (Map.lookup key vars) ++ drop 1 rest
                            where (key,rest) = break (==')') str
    interp1 plain = plain

sshblobFromPacket k = blob
 where
    Just (RSAKey (MPI n) (MPI e)) = rsaKeyFromPacket k
    bs = SSH.keyblob (n,e)
    blob = Char8.unpack bs

replaceSshServerKeys root cmn = do
    let homepass' = cmn { cap_homespec = fmap root (cap_homespec cmn) }
        replaceSSH op = op { opFiles = files }
            where
                files = Map.adjust delssh HomeSec
                        $ Map.adjust delssh HomePub
                        $ Map.insert (ArgFile $ root "/etc/ssh/ssh_host_rsa_key") strm $ opFiles op
                strm = streaminfo { typ = PEMFile, spill = KF_Match "ssh-server", access = Sec }
                delssh strm = strm { transforms = DeleteSubkeyByUsage "ssh-server" : transforms strm
                                   , fill = KF_All }
    KikiResult r report <- runKeyRing $ minimalOp homepass'
    case r of
        KikiSuccess rt -> Kiki.refreshCache rt $ case root "" of
                                                    "/" -> Nothing
                                                    ""  -> Nothing
                                                    pth -> Just pth
        err            -> hPutStrLn stderr $ errorString err

slash :: String -> String -> String
slash "/"    ('/':xs) = '/':xs
slash ""     ('/':xs) = '/':xs
slash ""     xs       = '/':xs
slash (y:ys) xs       = y:slash ys xs

ㄧchroot :: Args (FilePath -> FilePath)
ㄧchroot = pure (\r a -> slash r a) <*> arg "--chroot" <|> pure id

ㄧhomedir :: Args CommonArgsParsed
ㄧhomedir = CommonArgsParsed
                <$> optional (arg "--homedir")
                <*> optional (FileDesc <$> read <$> arg "--passphrase-fd")

kikiOptions :: ( [(String,Int)], [String] )
kikiOptions = ( ss, ps )
 where
    ss = [("--chroot",1),("--passphrase-fd",1),("--homedir",1)]
    ps = []