-- This file is part of Diohsc
-- Copyright (C) 2020 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 OverloadedStrings #-}

module ClientCert where

import           Control.Monad             ((<=<))
import           Control.Monad.Trans       (lift)
import           Control.Monad.Trans.Maybe (MaybeT (..), runMaybeT)
import           Crypto.Hash.Algorithms    (SHA256 (..))
import           Crypto.PubKey.RSA
import           Data.ASN1.BinaryEncoding  (DER (..))
import           Data.ASN1.Encoding        (decodeASN1', encodeASN1')
import           Data.ASN1.OID
import           Data.ASN1.Types           (ASN1Object (..))
import           Data.ASN1.Types.String    (ASN1StringEncoding (UTF8))
import           Data.Either               (fromRight)
import           Data.Foldable             (msum)
import           Data.Hourglass
import           Data.PEM
import           Data.X509
import           Safe
import           System.FilePath

import qualified Crypto.PubKey.Ed25519     as Ed25519
import qualified Crypto.PubKey.RSA.PKCS15  as PKCS15
import qualified Data.ByteArray            as BA
import qualified Data.ByteString           as BS
import qualified Data.Text                 as TS
import qualified Data.Text.Encoding        as TS

import           Fingerprint
import           Mundanities
import           Util

#ifndef WINDOWS
import           System.Posix.Files
#endif

-- |Certificate chain with secret key for tail cert
data ClientCert = ClientCert CertificateChain PrivKey
    deriving (Eq,Show)

clientCertFingerprint :: ClientCert -> Fingerprint
clientCertFingerprint (ClientCert (CertificateChain chain) _) =
    fingerprint $ head chain

maybeRight :: Either a b -> Maybe b
maybeRight (Left _)  = Nothing
maybeRight (Right b) = Just b

loadClientCert :: FilePath -> String -> IO (Maybe ClientCert)
loadClientCert path name =
    let certpath = path </> name <.> "crt"
        legacyKeypath = path </> name <.> "rsa"
        keypath = path </> name <.> "key"
    in ignoreIOErrAlt . runMaybeT $ do
        chain <- MaybeT $
            (maybeRight . decodeCertificateChain . CertificateChainRaw . map pemContent
                <=< maybeRight . pemParseBS)
            <$> BS.readFile certpath
        key <- msum
            [ MaybeT $ (maybeRight . (fst <$>) . fromASN1
                    <=< maybeRight . decodeASN1' DER . pemContent
                    <=< headMay
                    <=< maybeRight . pemParseBS)
                <$> ignoreIOErr (BS.readFile keypath)
            , do
                -- Legacy private rsa key format: Show instance of PrivateKey
                key <- MaybeT $ (PrivKeyRSA <$>) . readMay <$> ignoreIOErr (readFile legacyKeypath)
                -- Upgrade to standard format
                lift . saveClientCert path name $ ClientCert chain key
                return key
            ]
        return $ ClientCert chain key

saveClientCert :: FilePath -> String -> ClientCert -> IO ()
saveClientCert path name (ClientCert chain key) =
    let filepath = path </> name
        certpath = filepath <.> "crt"
        keypath = filepath <.> "key"
    in isSubPath path filepath >>? ignoreIOErr $ do
        let CertificateChainRaw rawCerts = encodeCertificateChain chain
            chainPEMs = map (pemWriteBS . PEM "CERTIFICATE" []) rawCerts
        BS.writeFile certpath $ BS.intercalate "\n" chainPEMs
        let header = case key of
                -- Use the header string the openssl commandline tool expects:
                PrivKeyRSA _ -> "RSA PRIVATE KEY"
                _            -> "PRIVATE KEY"
        BS.writeFile keypath . pemWriteBS . PEM header [] . encodeDER $ key
#ifndef WINDOWS
        setFileMode keypath $ unionFileModes ownerReadMode ownerWriteMode -- chmod 600
#endif
    where
    encodeDER :: ASN1Object o => o -> BS.ByteString
    encodeDER = encodeASN1' DER . (`toASN1` [])

-- RFC5280: To indicate that a certificate has no well-defined expiration
-- date, the notAfter SHOULD be assigned the GeneralizedTime value of
-- 99991231235959Z.
notAfterMax :: DateTime
notAfterMax = DateTime (Date 9999 December 31) (TimeOfDay 23 59 59 0)

-- RFC5280 has no corresponding prescription for notBefore, but
-- 19500101000000Z seems the canonical choice.
notBeforeMin :: DateTime
notBeforeMin = DateTime (Date 1950 January 1) (TimeOfDay 0 0 0 0)

data KeyType = KeyRSA | KeyEd25519

generateSelfSigned :: KeyType -> String -> IO ClientCert
generateSelfSigned tp cn =
    let dn = DistinguishedName [(getObjectID DnCommonName,
            ASN1CharacterString UTF8 . TS.encodeUtf8 $ TS.pack cn)]
        sigAlg = case tp of
            KeyRSA     -> SignatureALG HashSHA256 PubKeyALG_RSA
            KeyEd25519 -> SignatureALG_IntrinsicHash PubKeyALG_Ed25519
        to = timeConvert notAfterMax
        from = timeConvert notBeforeMin
        cert pubKey = Certificate
            { certVersion = 2
            , certSerial = 0
            , certSignatureAlg = sigAlg
            , certIssuerDN = dn
            , certSubjectDN = dn
            , certValidity = (from, to)
            , certPubKey = pubKey
            , certExtensions = Extensions Nothing
            }
    in case tp of
    KeyRSA -> do
        -- generate 2048bit RSA self-signed cert with maximum validity
        (pubKey, secKey) <- generate 256 65537
        blinder <- generateBlinder $ public_n pubKey
        let signed = fst $ objectToSignedExact
                (\b -> (fromRight BS.empty $
                    PKCS15.sign (Just blinder) (Just SHA256) secKey b, sigAlg, ()))
                (cert $ PubKeyRSA pubKey)
        return $ ClientCert (CertificateChain [signed]) (PrivKeyRSA secKey)
    KeyEd25519 -> do
        secKey <- Ed25519.generateSecretKey
        let pubKey = Ed25519.toPublic secKey
        let signed = fst $ objectToSignedExact
                (\b -> (BS.pack . BA.unpack $ Ed25519.sign secKey pubKey b, sigAlg, ()))
                (cert $ PubKeyEd25519 pubKey)
        return $ ClientCert (CertificateChain [signed]) (PrivKeyEd25519 secKey)