-- This file is part of htalkat
-- Copyright (C) 2021 Martin Bays <mbays@sdf.org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of version 3 of the GNU General Public License as
-- published by the Free Software Foundation, or any later version.
--
-- You should have received a copy of the GNU General Public License
-- along with this program.  If not, see http://www.gnu.org/licenses/.

{-# LANGUAGE CPP               #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}

module TLSTalk
    ( SocksProxy(..)
    , answerLast
    , connect
    , serve
    , spawnDefaultInteractiveClient
    ) where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad            (void, (<=<))
import           Data.Default.Class       (def)
import           Data.X509.Validation     (FailedReason (..), defaultChecks,
                                           defaultHooks, validate)
import           Network.Simple.TCP       (connectSock, connectSockSOCKS5)
import           Network.TLS              as TLS
import           Network.TLS.Extra.Cipher
import           System.Directory         (createDirectoryIfMissing)
import           System.Exit              (exitFailure)
import           System.FileLock          (SharedExclusive (..), withFileLock,
                                           withTryFileLock)
import           System.FilePath
import           System.IO                (IOMode (..), openFile)
import           System.IO.Temp           (withSystemTempDirectory)
import           System.Process           (rawSystem)
import           Time.System              (timeCurrent)
import           Time.Types               (Elapsed (..), Seconds (..))

import qualified Data.ByteString          as BS

import qualified Data.Text                as TS
import qualified Data.Text.Encoding       as TS
import qualified Data.X509                as X
import qualified Network.Simple.TCP       as TCP
import qualified Network.Socket           as S

import           Certificate
import           Config
import           Fingerprint
import           Host
import           Incoming
import           LookupPetname
import           Notify
import           Petname
import           RelayStream

#ifdef CURSES
import           CursesClient
#else
import           DumbClient
#endif

bindingNamedSocket :: FilePath -> (S.Socket -> IO a) -> IO a
bindingNamedSocket path =
    (`bracket` S.close) $ do
        sock <- S.socket S.AF_UNIX S.Stream 0
        S.bind sock $ S.SockAddrUnix path
        S.listen sock 1
        pure sock

serve :: FilePath -> Credential -> IO ()
serve ddir cred = errorOnNoLock <=< withTryFileLock listenLockPath Exclusive $ \_ -> do
    conf <- loadConfig ddir
    let params = def
                { serverShared = def { sharedCredentials = Credentials [cred] }
                , serverSupported = def
                    { supportedCiphers = talkatCiphersuite
                    , supportedVersions = [TLS13] }
                , serverHooks = def
                    { onUnverifiedClientCert = pure True
                    , onClientCertificate = \_ -> pure CertificateUsageAccept
                    }
                , serverWantClientCert = True
                }
    let port = show $ listen_port conf
        hostPref = case listen_host conf of
            [] -> TCP.HostAny
            h  -> TCP.Host h
        Just fp = spkiFingerprint <$> takeTailCert (fst cred)
    putStrLn $ "Listening on "
        <> (case hostPref of
            TCP.Host h -> "host " <> h
            _          -> "all available hosts")
        <> ", port " <> port <> "."
    putStrLn $ "URI: talkat:"
        <> showFingerprint fp <> "@"
        <> (let portBit = if port /= show defaultTalkatPort
                    then ":" <> port else ""
            in case hostPref of
                TCP.Host h -> h <> portBit
                _          -> "[host]" <> portBit <>
                    "\n       (replace [host] with your hostname or IP address)")
    cleanAllIncoming ddir
    serialMVar <- newMVar 0
    TCP.serve hostPref port $ \(sock,_) -> do
        S.setSocketOption sock S.NoDelay 1
        context <- contextNew sock params
        handshake context
        Just cert <- (takeTailCert =<<) <$> getClientCertificateChain context
        (if accept_unnamed conf then ((Just <$>) .) . lookupOrAddPetname else lookupPetname)
                ddir (spkiFingerprint cert) >>= \case
            Nothing -> pure ()
            Just petname -> withSystemTempDirectory "htalkat" $ \tdir -> do
                let sockPath = tdir </> "sock"
                bindingNamedSocket sockPath $ \dSock -> do
                    -- Serial numbers ensure we don't delete the wrong dir
                    serial <- modifyMVar serialMVar $ \n -> pure (n+1,n)
                    incoming <- addIncoming ddir cert sockPath serial
                    notifyOfIncoming ddir cert petname
                    relayStream context WriteFirst dSock
                    withFileLock (incomingDir ddir </> ".lock") Exclusive $ \_ ->
                        cleanIncoming ddir (Just serial) incoming
    where
    listenLockPath = ddir </> ".listen_lock"
    errorOnNoLock :: Maybe a -> IO ()
    errorOnNoLock Nothing = do
        putStrLn $ "Error: " <> listenLockPath <> " locked by another process."
        exitFailure
    errorOnNoLock _ = pure ()

data SocksProxy
    = NoSocksProxy
    | Socks5Proxy String String

connect :: FilePath -> Credential -> String -> SocksProxy -> Host -> Fingerprint -> IO ()
connect ddir cred name socksProxy (Host hostname port) fp = do
    let serverId = if port == defaultTalkatPort
            then BS.empty
            else TS.encodeUtf8 . TS.pack . (':':) $ show port
        params = (TLS.defaultParamsClient hostname serverId)
            { clientSupported = def
                { supportedCiphers = talkatCiphersuite
                , supportedVersions = [TLS13] }
            , clientHooks = def
                { onServerCertificate = checkServerCert
                , onCertificateRequest = \(_,_,_) ->
                    pure $ Just cred
                }
            }
    context <- do
        sock <- openSocket
        S.setSocketOption sock S.NoDelay 1
        c <- TLS.contextNew sock params
        handshake c >> pure c
    withSystemTempDirectory "htalkat" $ \tdir -> do
        let path = tdir </> "sock"
        bindingNamedSocket path $ \dSock -> do
            _ <- forkIO $ relayStream context WriteSecond dSock
            conf <- loadConfig ddir
            spawnInteractiveClient ddir conf name path
    where
    openSocket :: IO S.Socket
    openSocket = case socksProxy of
        NoSocksProxy -> fst <$> connectSock hostname (show port)
        Socks5Proxy socksHostname socksPort -> do
            sock <- fst <$> connectSock socksHostname socksPort
            _ <- connectSockSOCKS5 sock hostname (show port)
            pure sock

    checkServerCert store cache service chain | Just cert <- takeTailCert chain = do
        errors <- filter (not . ignoreError) <$> validate X.HashSHA256 defaultHooks
            (defaultChecks { checkExhaustive = True , checkLeafV3 = False }) store cache service chain
        if fp == spkiFingerprint cert then pure errors
            else do
                putStrLn "Server provides an unexpected certificate!"
                putStrLn $ "Expected: " <> showFingerprint fp
                putStrLn $ "Received: " <> showFingerprint (spkiFingerprint cert)
                pure $ UnknownCA : errors
        where
        ignoreError UnknownCA        = True
        ignoreError SelfSigned       = True
        ignoreError NotAnAuthority   = True
        ignoreError (NameMismatch _) = True
        ignoreError _                = False
    checkServerCert _ _ _ _ = pure [ EmptyChain ]

talkatCiphersuite :: [Cipher]
talkatCiphersuite =
    [ cipher_TLS13_AES128GCM_SHA256
    , cipher_TLS13_AES256GCM_SHA384
    , cipher_TLS13_CHACHA20POLY1305_SHA256
    , cipher_TLS13_AES128CCM_SHA256
    ]

answerLast :: FilePath -> Maybe Fingerprint -> IO ()
answerLast ddir mFp = do
    mInfo <- withFileLock (incomingDir ddir </> ".lock") Exclusive $ \_ -> do
        lastIncoming ddir mFp >>= \case
            Just incoming -> do
                petname <- incomingPetname ddir incoming
                sockPath <- readFile (incomingPath ddir incoming </> "sock")
                cleanIncoming ddir Nothing incoming
                pure $ Just (petname, sockPath)
            Nothing -> pure Nothing
    case mInfo of
        Just (petname, sockPath) -> do
            conf <- loadConfig ddir
            spawnInteractiveClient ddir conf (showPetname petname) sockPath
        Nothing -> putStrLn "Nothing to answer."

spawnInteractiveClient :: FilePath -> Config -> String -> FilePath -> IO ()
spawnInteractiveClient ddir conf name sockPath
    | command:args <- interactive_client conf =
        void . rawSystem command $ args ++ [name, sockPath]
    | otherwise =
#ifdef CURSES
        do
            mLog <- if curses_log conf
                then do
                    createDirectoryIfMissing True $ ddir </> "logs"
                    Elapsed (Seconds epochSecs) <- timeCurrent
                    Just <$> openFile (ddir </> "logs" </> name <>
                        "-" <> show epochSecs <.> "log") AppendMode
                else pure Nothing
            cursesClient (curses_local_top conf) mLog name sockPath
#else
        dumbClient sockPath
#endif

spawnDefaultInteractiveClient :: FilePath -> String -> FilePath -> IO ()
spawnDefaultInteractiveClient ddir name sockPath = do
    conf <- loadConfig ddir
    spawnInteractiveClient ddir (conf { interactive_client = [] })
        name sockPath