-- InjectSSHAgent.hs: hOpenPGP key tool inject-ssh-agent subcommand
-- Copyright © 2013-2026  Clint Adams
--
-- vim: softtabstop=4:shiftwidth=4:expandtab
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as
-- published by the Free Software Foundation, either version 3 of the
-- License, or (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.
{-# LANGUAGE GADTs #-}

module HOpenPGP.Tools.Hokey.InjectSSHAgent
    ( doInjectSSHAgent
    ) where

import Codec.Encryption.OpenPGP.Fingerprint (fingerprint)
import Codec.Encryption.OpenPGP.Ontology
    ( isKUF
    , isSKBindingSig
    )
import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Types
import Control.Exception (bracket)
import qualified Crypto.PubKey.RSA as RSA
import Data.Binary (get)
import Data.Binary.Get (getWord32be, runGet)
import Data.Binary.Put
    ( Put
    , putByteString
    , putLazyByteString
    , putWord32be
    , putWord8
    , runPut
    )
import Data.Bits (shiftR, testBit, (.&.))
import qualified Data.ByteString as B
import qualified Data.ByteString.Base16 as Base16
import qualified Data.ByteString.Char8 as BC8
import qualified Data.ByteString.Lazy as BL
import Data.Conduit (runConduitRes, (.|))
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Keyring
    ( AuthSecretSubkeyAtTime
    , authSecretSubkeyPrimaryUID
    , authSecretSubkeyValue
    , conduitToAuthSecretSubkeysAt
    , conduitToSomeTKsDroppingEither
    )
import Data.Conduit.Serialization.Binary (conduitGet)
import Data.Foldable (find)
import Data.List (intercalate, sortOn)
import Data.Maybe (fromMaybe, listToMaybe, mapMaybe)
import Data.Ord (Down (..))
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as T
import Data.Time.Clock.POSIX
    ( POSIXTime
    , getPOSIXTime
    , posixSecondsToUTCTime
    )
import Network.Socket
    ( Family (AF_UNIX)
    , SockAddr (..)
    , Socket
    , SocketType (Stream)
    , close
    , connect
    , defaultProtocol
    , socket
    )
import qualified Network.Socket.ByteString as NSB
import System.Environment (lookupEnv)
import System.Exit (exitFailure)
import System.IO
    ( hPutStrLn
    , stderr
    , stdin
    )

import HOpenPGP.Tools.Common.Common (renderFingerprint)
import HOpenPGP.Tools.Hokey.Options (InjectSSHAgentOptions (..))

doInjectSSHAgent :: InjectSSHAgentOptions -> IO ()
doInjectSSHAgent opts = do
    socketPath <-
        resolveSSHAgentSocketPath (injectSSHAgentSocket opts)
    input <- readInjectedSecretKeyMaterial opts
    cpt <- getPOSIXTime
    authCandidates <-
        runConduitRes $
            CL.sourceList (BL.toChunks input)
                .| conduitGet get
                .| conduitToSomeTKsDroppingEither
                .| CL.mapMaybe (either (const Nothing) (>>= someTKToSecretTK))
                .| conduitToAuthSecretSubkeysAt (posixSecondsToUTCTime cpt)
                .| CL.consume
    injectableCandidates <-
        if null authCandidates
            then inferInjectableAuthSubkeys cpt input
            else pure (map candidateFromAuthSecretSubkey authCandidates)
    whenEmpty
        injectableCandidates
        "inject-ssh-agent: no authentication-capable secret subkey found"
    selectedRequests <-
        selectInjectableAuthSubkeys
            (injectSSHAgentComment opts)
            injectableCandidates
    mapM_
        ( \(selected, request) -> do
            sendAddIdentityToSSHAgent socketPath request
            hPutStrLn stderr $
                "inject-ssh-agent: added authentication subkey "
                    ++ renderFingerprint
                        (fingerprint (injectableAuthSubkeyPKP selected))
                    ++ " to ssh-agent"
        )
        selectedRequests

candidateFromAuthSecretSubkey
    :: AuthSecretSubkeyAtTime -> InjectableAuthSubkey
candidateFromAuthSecretSubkey authSubkey =
    InjectableAuthSubkey
        { injectableAuthSubkeyPKP = authSecretSubkeyPKP authSubkey
        , injectableAuthSubkeySKA = authSecretSubkeySKA authSubkey
        , injectableAuthSubkeyPrimaryUID =
            authSecretSubkeyPrimaryUID authSubkey
        }

inferInjectableAuthSubkeys
    :: POSIXTime -> BL.ByteString -> IO [InjectableAuthSubkey]
inferInjectableAuthSubkeys _cpt input = do
    tks <-
        runConduitRes $
            CL.sourceList (BL.toChunks input)
                .| conduitGet get
                .| conduitToSomeTKsDroppingEither
                .| CL.mapMaybe (either (const Nothing) (>>= someTKToSecretTK))
                .| CL.consume
    pure (concatMap inferFromTK tks)
  where
    inferFromTK tk =
        let mPrimaryUID = fst <$> listToMaybe (_tkUIDs tk)
         in mapMaybe (inferFromSubkey mPrimaryUID) (_tkSubs tk)
    inferFromSubkey
        :: Maybe Text
        -> (KeyPkt k, [SignaturePayload])
        -> Maybe InjectableAuthSubkey
    inferFromSubkey mPrimaryUID (KeyPktSecretSubkey pkp ska, sigs)
        | hasAuthCapability sigs =
            Just
                InjectableAuthSubkey
                    { injectableAuthSubkeyPKP = pkp
                    , injectableAuthSubkeySKA = ska
                    , injectableAuthSubkeyPrimaryUID = mPrimaryUID
                    }
    inferFromSubkey _ _ = Nothing
    hasAuthCapability sigs =
        any
            (Set.member AuthKey)
            ( mapMaybe
                signatureKeyFlags
                (newestWithUsageFlags (filter isSKBindingSig sigs))
            )
    newestWithUsageFlags =
        take 1
            . sortOn (Down . take 1 . sigCreationTimes)
            . filter (any isKUF . signatureHashedSubpackets)
    sigCreationTimes =
        mapMaybe sigCreationTimeFromSubpacket . signatureHashedSubpackets
    sigCreationTimeFromSubpacket (SigSubPacket _ (SigCreationTime ct)) = Just ct
    sigCreationTimeFromSubpacket _ = Nothing
    signatureHashedSubpackets (SigV4 _ _ _ hasheds _ _ _) = hasheds
    signatureHashedSubpackets (SigV6 _ _ _ _ hasheds _ _ _) = hasheds
    signatureHashedSubpackets _ = []
    signatureKeyFlags sig = do
        sp <- find isKUF (signatureHashedSubpackets sig)
        case sp of
            SigSubPacket _ (KeyFlags flags) -> Just flags
            _ -> Nothing

resolveSSHAgentSocketPath :: Maybe String -> IO String
resolveSSHAgentSocketPath (Just path) = pure path
resolveSSHAgentSocketPath Nothing = do
    envPath <- lookupEnv "SSH_AUTH_SOCK"
    case envPath of
        Just path -> pure path
        Nothing ->
            failInject
                "inject-ssh-agent: SSH_AUTH_SOCK is not set; use --ssh-agent-socket"

readInjectedSecretKeyMaterial
    :: InjectSSHAgentOptions -> IO BL.ByteString
readInjectedSecretKeyMaterial opts = do
    chunks <-
        case injectSSHAgentFromFD opts of
            Nothing -> runConduitRes $ CB.sourceHandle stdin .| CL.consume
            Just fd
                | fd < 0 ->
                    failInject
                        "inject-ssh-agent: --from-fd must be a non-negative integer"
                | otherwise ->
                    runConduitRes $
                        CB.sourceFile ("/dev/fd/" ++ show fd) .| CL.consume
    let input = BL.fromChunks chunks
    if BL.null input
        then
            failInject
                "inject-ssh-agent: no secret key bytes were provided on the selected input stream"
        else pure input

selectInjectableAuthSubkeys
    :: Maybe String
    -> [InjectableAuthSubkey]
    -> IO [(InjectableAuthSubkey, BL.ByteString)]
selectInjectableAuthSubkeys mComment candidates
    | null selectedRequests =
        failInject
            ( "inject-ssh-agent: auth-capable subkeys were found, but none are supported for ssh-agent injection: "
                ++ intercalate "; " (reverse errs)
            )
    | otherwise = pure (reverse selectedRequests)
  where
    (errs, selectedRequests) = foldl' pick ([], []) candidates
    pick (accErrs, accSelected) candidate =
        case sshAddIdentityRequest
            (fromMaybe (defaultSSHComment candidate) mComment)
            (injectableAuthSubkeyPKP candidate)
            (injectableAuthSubkeySKA candidate) of
            Left err -> (err : accErrs, accSelected)
            Right request -> (accErrs, (candidate, request) : accSelected)
    defaultSSHComment candidate =
        case injectableAuthSubkeyPrimaryUID candidate of
            Just uid -> T.unpack uid
            Nothing ->
                "openpgp:"
                    ++ BC8.unpack
                        ( Base16.encode
                            ( BL.toStrict
                                (unFingerprint (fingerprint (injectableAuthSubkeyPKP candidate)))
                            )
                        )

sshAddIdentityRequest
    :: String
    -> SomePKPayload
    -> SKAddendum
    -> Either String BL.ByteString
sshAddIdentityRequest comment subkeyPKP subkeySKA =
    case subkeySKA of
        SUUnencrypted (RSAPrivateKey (RSA_PrivateKey rsaPrivateKey)) _ ->
            Right $
                frameSSHAgentRequest
                    (rsaAddIdentityPayload (BC8.pack comment) rsaPrivateKey)
        SUUnencrypted (EdDSAPrivateKey EdSigningCurve25519 secretSeed) _ ->
            frameSSHAgentRequest
                <$> ed25519AddIdentityPayload (BC8.pack comment) subkeyPKP secretSeed
        SUUnencrypted (UnknownSKey rawSecret) _
            | isEd25519PKA (_pkalgo subkeyPKP) ->
                frameSSHAgentRequest
                    <$> ed25519AddIdentityPayload
                        (BC8.pack comment)
                        subkeyPKP
                        (BL.toStrict rawSecret)
        SUUnencrypted (EdDSAPrivateKey EdSigningCurve448 _) _ ->
            Left
                ( "subkey "
                    ++ renderFingerprint (fingerprint subkeyPKP)
                    ++ " uses Ed448, which is not supported by ssh-agent add-identity"
                )
        SUUnencrypted _ _ ->
            Left
                ( "subkey "
                    ++ renderFingerprint (fingerprint subkeyPKP)
                    ++ " uses an unsupported key algorithm for ssh-agent injection"
                )
        _ ->
            Left
                ( "subkey "
                    ++ renderFingerprint (fingerprint subkeyPKP)
                    ++ " is encrypted; decrypt it before injection"
                )

authSecretSubkeyPKP :: AuthSecretSubkeyAtTime -> SomePKPayload
authSecretSubkeyPKP = keyPktPKPayload . authSecretSubkeyValue

authSecretSubkeySKA :: AuthSecretSubkeyAtTime -> SKAddendum
authSecretSubkeySKA = secretKeyPktSKAddendum . authSecretSubkeyValue

rsaAddIdentityPayload
    :: B.ByteString -> RSA.PrivateKey -> BL.ByteString
rsaAddIdentityPayload comment privateKey =
    runPut $ do
        putWord8 17
        putSSHString (BC8.pack "ssh-rsa")
        putSSHMpint (RSA.public_n (RSA.private_pub privateKey))
        putSSHMpint (RSA.public_e (RSA.private_pub privateKey))
        putSSHMpint (RSA.private_d privateKey)
        putSSHMpint (RSA.private_qinv privateKey)
        putSSHMpint (RSA.private_p privateKey)
        putSSHMpint (RSA.private_q privateKey)
        putSSHString comment

ed25519AddIdentityPayload
    :: B.ByteString
    -> SomePKPayload
    -> B.ByteString
    -> Either String BL.ByteString
ed25519AddIdentityPayload comment pkp rawSecret = do
    publicKey <- ed25519PublicPoint pkp
    secretSeed <- normalizeEd25519Secret rawSecret
    pure $
        runPut $ do
            putWord8 17
            putSSHString (BC8.pack "ssh-ed25519")
            putSSHString publicKey
            putSSHString (secretSeed <> publicKey)
            putSSHString comment

ed25519PublicPoint :: SomePKPayload -> Either String B.ByteString
ed25519PublicPoint pkp =
    case _pubkey pkp of
        EdDSAPubKey EdSigningCurve25519 point ->
            maybe
                ( Left
                    ( "invalid Ed25519 public point for subkey "
                        ++ renderFingerprint (fingerprint pkp)
                    )
                )
                Right
                (edPointToRawBytes point)
        _ ->
            Left
                ( "subkey "
                    ++ renderFingerprint (fingerprint pkp)
                    ++ " does not have an Ed25519 public key"
                )

edPointToRawBytes :: EdPoint -> Maybe B.ByteString
edPointToRawBytes (NativeEPoint (EPoint i)) = integerToFixedBytes 32 i
edPointToRawBytes (PrefixedNativeEPoint (EPoint i)) = do
    prefixed <- integerToFixedBytes 33 i
    case B.uncons prefixed of
        Just (0x40, raw) -> Just raw
        _ -> Nothing

normalizeEd25519Secret
    :: B.ByteString -> Either String B.ByteString
normalizeEd25519Secret rawSecret
    | B.length rawSecret == 32 = Right rawSecret
    | otherwise =
        Left
            ( "expected 32-byte Ed25519 secret seed, got "
                ++ show (B.length rawSecret)
                ++ " bytes"
            )

putSSHString :: B.ByteString -> Put
putSSHString bs = putWord32be (fromIntegral (B.length bs)) >> putByteString bs

putSSHMpint :: Integer -> Put
putSSHMpint n
    | n <= 0 = putWord32be 0
    | otherwise = putSSHString encoded
  where
    raw = integerToUnsignedBytes n
    encoded =
        case B.uncons raw of
            Just (firstByte, _)
                | testBit firstByte 7 -> B.cons 0x00 raw
            _ -> raw

integerToFixedBytes :: Int -> Integer -> Maybe B.ByteString
integerToFixedBytes width n
    | n < 0 = Nothing
    | B.length raw > width = Nothing
    | otherwise =
        Just (B.replicate (width - B.length raw) 0x00 <> raw)
  where
    raw =
        if n == 0
            then B.singleton 0x00
            else integerToUnsignedBytes n

integerToUnsignedBytes :: Integer -> B.ByteString
integerToUnsignedBytes n =
    B.reverse $
        B.unfoldr
            ( \value ->
                if value == 0
                    then Nothing
                    else Just (fromIntegral (value .&. 0xff), value `shiftR` 8)
            )
            n

isEd25519PKA :: PubKeyAlgorithm -> Bool
isEd25519PKA pka = fromFVal pka == 27

frameSSHAgentRequest :: BL.ByteString -> BL.ByteString
frameSSHAgentRequest body =
    runPut $
        putWord32be (fromIntegral (BL.length body))
            >> putLazyByteString body

sendAddIdentityToSSHAgent :: FilePath -> BL.ByteString -> IO ()
sendAddIdentityToSSHAgent socketPath request =
    bracket
        (socket AF_UNIX Stream defaultProtocol)
        close
        ( \sock -> do
            connect sock (SockAddrUnix socketPath)
            NSB.sendAll sock (BL.toStrict request)
            response <- readSSHAgentPacket sock
            case B.uncons response of
                Just (6, _) -> pure ()
                Just (5, _) ->
                    failInject
                        "inject-ssh-agent: ssh-agent rejected the supplied key"
                Just (code, _) ->
                    failInject
                        ( "inject-ssh-agent: ssh-agent returned unexpected response type "
                            ++ show code
                        )
                Nothing ->
                    failInject
                        "inject-ssh-agent: ssh-agent returned an empty response packet"
        )

readSSHAgentPacket :: Socket -> IO B.ByteString
readSSHAgentPacket sock = do
    lenPrefix <- recvExact sock 4
    let packetLen = fromIntegral (runGet getWord32be (BL.fromStrict lenPrefix))
    recvExact sock packetLen

recvExact :: Socket -> Int -> IO B.ByteString
recvExact _ 0 = pure B.empty
recvExact sock remaining = go B.empty remaining
  where
    go acc 0 = pure acc
    go acc bytesRemaining = do
        chunk <- NSB.recv sock bytesRemaining
        if B.null chunk
            then
                failInject
                    "inject-ssh-agent: ssh-agent socket closed while reading response"
            else go (acc <> chunk) (bytesRemaining - B.length chunk)

whenEmpty :: [a] -> String -> IO ()
whenEmpty [] msg = failInject msg
whenEmpty _ _ = pure ()

failInject :: String -> IO a
failInject msg = hPutStrLn stderr msg >> exitFailure

data InjectableAuthSubkey
    = InjectableAuthSubkey
    { injectableAuthSubkeyPKP :: SomePKPayload
    , injectableAuthSubkeySKA :: SKAddendum
    , injectableAuthSubkeyPrimaryUID :: Maybe Text
    }
