{-# LANGUAGE DeriveDataTypeable, TypeFamilies #-}

module Propellor.Property.Ssh (
        installed,
        restarted,
        PubKeyText,
        SshKeyType(..),
        -- * Daemon configuration
        sshdConfig,
        ConfigKeyword,
        setSshdConfigBool,
        setSshdConfig,
        RootLogin(..),
        permitRootLogin,
        passwordAuthentication,
        noPasswords,
        listenPort,
        -- * Host keys
        randomHostKeys,
        hostKeys,
        hostKey,
        hostPubKey,
        getHostPubKey,
        -- * User keys and configuration
        userKeys,
        userKeyAt,
        knownHost,
        unknownHost,
        authorizedKeysFrom,
        unauthorizedKeysFrom,
        authorizedKeys,
        authorizedKey,
        hasAuthorizedKeys,
        getUserPubKeys,
) where

import Propellor.Base
import qualified Propellor.Property.File as File
import qualified Propellor.Property.Service as Service
import qualified Propellor.Property.Apt as Apt
import Propellor.Property.User
import Propellor.Types.Info
import Utility.FileMode

import System.PosixCompat
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.Semigroup as Sem
import Data.List

installed :: Property UnixLike
installed = "ssh installed" ==> (aptinstall `pickOS` unsupportedOS)
  where
        aptinstall :: Property DebianLike
        aptinstall = Apt.installed ["ssh"]

restarted :: Property DebianLike
restarted = Service.restarted "ssh"

sshBool :: Bool -> String
sshBool True = "yes"
sshBool False = "no"

sshdConfig :: FilePath
sshdConfig = "/etc/ssh/sshd_config"

type ConfigKeyword = String

setSshdConfigBool :: ConfigKeyword -> Bool -> Property DebianLike
setSshdConfigBool setting allowed = setSshdConfig setting (sshBool allowed)

setSshdConfig :: ConfigKeyword -> String -> Property DebianLike
setSshdConfig setting v = File.fileProperty desc f sshdConfig
        `onChange` restarted
  where
        desc = unwords [ "ssh config:", setting, v ]
        cfgline = setting ++ " " ++ v
        wantedline s
                | s == cfgline = True
                | (setting ++ " ") `isPrefixOf` s = False
                | otherwise = True
        f ls
                | cfgline `elem` ls = filter wantedline ls
                | otherwise = filter wantedline ls ++ [cfgline]

data RootLogin
        = RootLogin Bool  -- ^ allow or prevent root login
        | WithoutPassword -- ^ disable password authentication for root, while allowing other authentication methods
        | ForcedCommandsOnly -- ^ allow root login with public-key authentication, but only if a forced command has been specified for the public key

permitRootLogin :: RootLogin -> Property DebianLike
permitRootLogin (RootLogin b) = setSshdConfigBool "PermitRootLogin" b
permitRootLogin WithoutPassword = setSshdConfig "PermitRootLogin" "without-password"
permitRootLogin ForcedCommandsOnly = setSshdConfig "PermitRootLogin" "forced-commands-only"

passwordAuthentication :: Bool -> Property DebianLike
passwordAuthentication = setSshdConfigBool "PasswordAuthentication"

-- | Configure ssh to not allow password logins.
--
-- To prevent lock-out, this is done only once root's
-- authorized_keys is in place.
noPasswords :: Property DebianLike
noPasswords = check (hasAuthorizedKeys (User "root")) $
        passwordAuthentication False

dotDir :: User -> IO FilePath
dotDir user = do
        h <- homedir user
        return $ h </> ".ssh"

dotFile :: FilePath -> User -> IO FilePath
dotFile f user = do
        d <- dotDir user
        return $ d </> f

-- | Makes the ssh server listen on a given port, in addition to any other
-- ports it is configured to listen on.
--
-- Revert to prevent it listening on a particular port.
listenPort :: Port -> RevertableProperty DebianLike DebianLike
listenPort port = enable <!> disable
  where
        portline = "Port " ++ val port
        enable = sshdConfig `File.containsLine` portline
                `describe` ("ssh listening on " ++ portline)
                `onChange` restarted
        disable = sshdConfig `File.lacksLine` portline
                `describe` ("ssh not listening on " ++ portline)
                `onChange` restarted

hasAuthorizedKeys :: User -> IO Bool
hasAuthorizedKeys = go <=< dotFile "authorized_keys"
  where
        go f = not . null <$> catchDefaultIO "" (readFile f)

-- | Blows away existing host keys and make new ones.
-- Useful for systems installed from an image that might reuse host keys.
-- A flag file is used to only ever do this once.
randomHostKeys :: Property DebianLike
randomHostKeys = flagFile prop "/etc/ssh/.unique_host_keys"
        `onChange` restarted
  where
        prop :: Property UnixLike
        prop = property' "ssh random host keys" $ \w -> do
                void $ liftIO $ boolSystem "sh"
                        [ Param "-c"
                        , Param "rm -f /etc/ssh/ssh_host_*"
                        ]
                ensureProperty w $ scriptProperty [ "DPKG_MAINTSCRIPT_NAME=postinst DPKG_MAINTSCRIPT_PACKAGE=openssh-server /var/lib/dpkg/info/openssh-server.postinst configure" ]
                        `assume` MadeChange

-- | The text of a ssh public key, for example, "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIB3BJ2GqZiTR2LEoDXyYFgh/BduWefjdKXAsAtzS9zeI"
type PubKeyText = String

-- | Installs the specified list of ssh host keys.
--
-- The corresponding private keys come from the privdata.
--
-- Any host keys that are not in the list are removed from the host.
hostKeys :: IsContext c => c -> [(SshKeyType, PubKeyText)] -> Property (HasInfo + DebianLike)
hostKeys ctx l = go `before` cleanup
  where
        desc = "ssh host keys configured " ++ typelist (map fst l)
        go :: Property (HasInfo + DebianLike)
        go = propertyList desc $ toProps $ catMaybes $
                map (\(t, pub) -> Just $ hostKey ctx t pub) l
        typelist tl = "(" ++ unwords (map fromKeyType tl) ++ ")"
        alltypes = [minBound..maxBound]
        staletypes = let have = map fst l in filter (`notElem` have) alltypes
        removestale :: Bool -> [Property DebianLike]
        removestale b = map (tightenTargets . File.notPresent . flip keyFile b) staletypes
        cleanup :: Property DebianLike
        cleanup
                | null staletypes || null l = doNothing
                | otherwise =
                        combineProperties ("any other ssh host keys removed " ++ typelist staletypes)
                                (toProps $ removestale True ++ removestale False)
                                `onChange` restarted

-- | Installs a single ssh host key of a particular type.
--
-- The public key is provided to this function;
-- the private key comes from the privdata;
hostKey :: IsContext c => c -> SshKeyType -> PubKeyText -> Property (HasInfo + DebianLike)
hostKey context keytype pub = go `onChange` restarted
  where
        go = combineProperties desc $ props
                & hostPubKey keytype pub
                & installpub
                & installpriv
        desc = "ssh host key configured (" ++ fromKeyType keytype ++ ")"
        keysrc ext field = PrivDataSourceFileFromCommand field ("sshkey"++ext)
                ("ssh-keygen -t " ++ sshKeyTypeParam keytype ++ " -f sshkey")
        installpub :: Property UnixLike
        installpub = keywriter File.hasContent True (lines pub)
        installpriv :: Property (HasInfo + UnixLike)
        installpriv = withPrivData (keysrc "" (SshPrivKey keytype "")) context $ \getkey ->
                property' desc $ \w -> getkey $
                        ensureProperty w
                                . keywriter File.hasContentProtected False
                                . privDataLines
        keywriter p ispub keylines = do
                let f = keyFile keytype ispub
                p f (keyFileContent keylines)

-- Make sure that there is a newline at the end;
-- ssh requires this for some types of private keys.
keyFileContent :: [String] -> [File.Line]
keyFileContent keylines = keylines ++ [""]

keyFile :: SshKeyType -> Bool -> FilePath
keyFile keytype ispub = "/etc/ssh/ssh_host_" ++ fromKeyType keytype ++ "_key" ++ ext
  where
        ext = if ispub then ".pub" else ""

-- | Indicates the host key that is used by a Host, but does not actually
-- configure the host to use it. Normally this does not need to be used;
-- use 'hostKey' instead.
hostPubKey :: SshKeyType -> PubKeyText -> Property (HasInfo + UnixLike)
hostPubKey t = pureInfoProperty "ssh pubkey known" . HostKeyInfo . M.singleton t

getHostPubKey :: Propellor (M.Map SshKeyType PubKeyText)
getHostPubKey = fromHostKeyInfo <$> askInfo

newtype HostKeyInfo = HostKeyInfo
        { fromHostKeyInfo :: M.Map SshKeyType PubKeyText }
        deriving (Eq, Ord, Typeable, Show)

instance IsInfo HostKeyInfo where
        propagateInfo _ = PropagateInfo False

instance Sem.Semigroup HostKeyInfo where
        HostKeyInfo old <> HostKeyInfo new =
                -- new first because union prefers values from the first
                -- parameter when there is a duplicate key
                HostKeyInfo (new `M.union` old)

instance Monoid HostKeyInfo where
        mempty = HostKeyInfo M.empty
        mappend = (Sem.<>)

userPubKeys :: User -> [(SshKeyType, PubKeyText)] -> Property (HasInfo + UnixLike)
userPubKeys u@(User n) l = pureInfoProperty ("ssh pubkey for " ++ n) $
        UserKeyInfo (M.singleton u (S.fromList l))

getUserPubKeys :: User -> Propellor [(SshKeyType, PubKeyText)]
getUserPubKeys u = maybe [] S.toList . M.lookup u . fromUserKeyInfo <$> askInfo

newtype UserKeyInfo = UserKeyInfo
        { fromUserKeyInfo :: M.Map User (S.Set (SshKeyType, PubKeyText)) }
        deriving (Eq, Ord, Typeable, Show)

instance IsInfo UserKeyInfo where
        propagateInfo _ = PropagateInfo False

instance Sem.Semigroup UserKeyInfo where
        UserKeyInfo old <> UserKeyInfo new =
                UserKeyInfo (M.unionWith S.union old new)

instance Monoid UserKeyInfo where
        mempty = UserKeyInfo M.empty
        mappend = (Sem.<>)

-- | Sets up a user with the specified public keys, and the corresponding
-- private keys from the privdata.
--
-- The public keys are added to the Info, so other properties like
-- `authorizedKeysFrom` can use them.
userKeys :: IsContext c => User -> c -> [(SshKeyType, PubKeyText)] -> Property (HasInfo + UnixLike)
userKeys user@(User name) context ks = combineProperties desc $ toProps $
        userPubKeys user ks : map (userKeyAt Nothing user context) ks
  where
        desc = unwords
                [ name
                , "has ssh key"
                , "(" ++ unwords (map (fromKeyType . fst) ks) ++ ")"
                ]

-- | Sets up a user with the specified pubic key, and a private
-- key from the privdata.
--
-- A file can be specified to write the key to somewhere other than
-- the default locations. Allows a user to have multiple keys for
-- different roles.
userKeyAt :: IsContext c => Maybe FilePath -> User -> c -> (SshKeyType, PubKeyText) -> Property (HasInfo + UnixLike)
userKeyAt dest user@(User u) context (keytype, pubkeytext) =
        combineProperties desc $ props
                & pubkey
                & privkey
  where
        desc = unwords $ catMaybes
                [ Just u
                , Just "has ssh key"
                , dest
                , Just $ "(" ++ fromKeyType keytype ++ ")"
                ]
        pubkey :: Property UnixLike
        pubkey = property' desc $ \w ->
                ensureProperty w =<< installprop File.hasContent ".pub" [pubkeytext]
        privkey :: Property (HasInfo + UnixLike)
        privkey = withPrivData (SshPrivKey keytype u) context privkey'
        privkey' :: ((PrivData -> Propellor Result) -> Propellor Result) -> Property (HasInfo + UnixLike)
        privkey' getkey = property' desc $ \w -> getkey $ \k ->
                ensureProperty w
                        =<< installprop File.hasContentProtected "" (privDataLines k)
        installprop writer ext key = do
                f <- liftIO $ keyfile ext
                return $ combineProperties desc $ props
                        & writer f (keyFileContent key)
                        & File.ownerGroup f user (userGroup user)
                        & File.ownerGroup (takeDirectory f) user (userGroup user)
        keyfile ext = case dest of
                Nothing -> do
                        home <- homeDirectory <$> getUserEntryForName u
                        return $ home </> ".ssh" </> "id_" ++ fromKeyType keytype ++ ext
                Just f -> return $ f ++ ext

fromKeyType :: SshKeyType -> String
fromKeyType SshRsa = "rsa"
fromKeyType SshDsa = "dsa"
fromKeyType SshEcdsa = "ecdsa"
fromKeyType SshEd25519 = "ed25519"

-- | Puts some host's ssh public key(s), as set using `hostPubKey`
-- or `hostKey` into the known_hosts file for a user.
knownHost :: [Host] -> HostName -> User -> Property UnixLike
knownHost hosts hn user@(User u) = property' desc $ \w ->
        go w =<< knownHostLines hosts hn
  where
        desc = u ++ " knows ssh key for " ++ hn

        go _ [] = do
                warningMessage $ "no configured ssh host keys for " ++ hn
                return FailedChange
        go w ls = do
                f <- liftIO $ dotFile "known_hosts" user
                ensureProperty w $ modKnownHost user f $
                        f `File.containsLines` ls
                                `requires` File.dirExists (takeDirectory f)

-- | Reverts `knownHost`
unknownHost :: [Host] -> HostName -> User -> Property UnixLike
unknownHost hosts hn user@(User u) = property' desc $ \w ->
        go w =<< knownHostLines hosts hn
  where
        desc = u ++ " does not know ssh key for " ++ hn

        go _ [] = return NoChange
        go w ls = do
                f <- liftIO $ dotFile "known_hosts" user
                ifM (liftIO $ doesFileExist f)
                        ( ensureProperty w $ modKnownHost user f $
                                f `File.lacksLines` ls
                        , return NoChange
                        )

knownHostLines :: [Host] -> HostName -> Propellor [File.Line]
knownHostLines hosts hn = keylines <$> fromHost hosts hn getHostPubKey
  where
        keylines (Just m) = map (\k -> hn ++ " " ++ k) (M.elems m)
        keylines Nothing = []

modKnownHost :: User -> FilePath -> Property UnixLike -> Property UnixLike
modKnownHost user f p = p
        `before` File.ownerGroup f user (userGroup user)
        `before` File.ownerGroup (takeDirectory f) user (userGroup user)

-- | Ensures that a local user's authorized_keys contains lines allowing
-- logins from a remote user on the specified Host.
--
-- The ssh keys of the remote user can be set using `userKeys`
--
-- Any other lines in the authorized_keys file are preserved as-is.
authorizedKeysFrom :: User -> (User, Host) -> Property UnixLike
localuser@(User ln) `authorizedKeysFrom` (remoteuser@(User rn), remotehost) =
        property' desc (\w -> go w =<< authorizedKeyLines remoteuser remotehost)
  where
        remote = rn ++ "@" ++ hostName remotehost
        desc = ln ++ " authorized_keys from " ++ remote

        go _ [] = do
                warningMessage $ "no configured ssh user keys for " ++ remote
                return FailedChange
        go w ls = ensureProperty w $ combineProperties desc $ toProps $
                map (setupRevertableProperty . authorizedKey localuser) ls

-- | Reverts `authorizedKeysFrom`
unauthorizedKeysFrom :: User -> (User, Host) -> Property UnixLike
localuser@(User ln) `unauthorizedKeysFrom` (remoteuser@(User rn), remotehost) =
        property' desc (\w -> go w =<< authorizedKeyLines remoteuser remotehost)
  where
        remote = rn ++ "@" ++ hostName remotehost
        desc = ln ++ " unauthorized_keys from " ++ remote

        go _ [] = return NoChange
        go w ls = ensureProperty w $ combineProperties desc $ toProps $
                map (undoRevertableProperty . authorizedKey localuser) ls

authorizedKeyLines :: User -> Host -> Propellor [File.Line]
authorizedKeyLines remoteuser remotehost =
        map snd <$> fromHost' remotehost (getUserPubKeys remoteuser)

-- | Makes a user have authorized_keys from the PrivData
--
-- This removes any other lines from the file.
authorizedKeys :: IsContext c => User -> c -> Property (HasInfo + UnixLike)
authorizedKeys user@(User u) context = withPrivData (SshAuthorizedKeys u) context $ \get ->
        property' desc $ \w -> get $ \v -> do
                f <- liftIO $ dotFile "authorized_keys" user
                ensureProperty w $ combineProperties desc $ props
                        & File.hasContentProtected f (keyFileContent (privDataLines v))
                        & File.ownerGroup f user (userGroup user)
                        & File.ownerGroup (takeDirectory f) user (userGroup user)
  where
        desc = u ++ " has authorized_keys"

-- | Ensures that a user's authorized_keys contains a line.
-- Any other lines in the file are preserved as-is.
authorizedKey :: User -> String -> RevertableProperty UnixLike UnixLike
authorizedKey user@(User u) l = add <!> remove
  where
        add = property' (u ++ " has authorized_keys") $ \w -> do
                f <- liftIO $ dotFile "authorized_keys" user
                ensureProperty w $ modAuthorizedKey f user $
                        f `File.containsLine` l
                                `requires` File.dirExists (takeDirectory f)
        remove = property' (u ++ " lacks authorized_keys") $ \w -> do
                f <- liftIO $ dotFile "authorized_keys" user
                ifM (liftIO $ doesFileExist f)
                        ( ensureProperty w $ modAuthorizedKey f user $
                                f `File.lacksLine` l
                        , return NoChange
                        )

modAuthorizedKey :: FilePath -> User -> Property UnixLike -> Property UnixLike
modAuthorizedKey f user p = p
        `before` File.mode f (combineModes [ownerWriteMode, ownerReadMode])
        `before` File.ownerGroup f user (userGroup user)
        `before` File.ownerGroup (takeDirectory f) user (userGroup user)