{-# LANGUAGE GADTs #-}

-- | Types and functions used to build test certificates.

module Certificate

    (

    -- * Hash algorithms

      hashMD2

    , hashMD5

    , hashSHA1

    , hashSHA224

    , hashSHA256

    , hashSHA384

    , hashSHA512

    -- * Key and signature utilities

    , Alg(..)

    , Keys

    , generateKeys

    -- * Certificate utilities

    , Pair(..)

    , mkDn

    , mkExtension

    , leafStdExts

    -- * Certificate creation functions

    , Auth(..)

    , mkCertificate

    , mkCA

    , mkLeaf

    ) where

import Control.Applicative

import Crypto.Hash.Algorithms

import Crypto.Number.Serialize

import qualified Crypto.PubKey.DSA        as DSA

import qualified Crypto.PubKey.ECC.ECDSA  as ECDSA

import qualified Crypto.PubKey.ECC.Generate as ECC

import qualified Crypto.PubKey.ECC.Types  as ECC

import qualified Crypto.PubKey.RSA        as RSA

import qualified Crypto.PubKey.RSA.PKCS15 as RSA

import qualified Crypto.PubKey.RSA.PSS    as PSS

import qualified Data.ByteString as B

import Data.ASN1.BinaryEncoding (DER(..))

import Data.ASN1.Encoding

import Data.ASN1.Types

import Data.Maybe (catMaybes)

import Data.String (fromString)

import Data.X509

import Data.Hourglass

-- Crypto utilities --

-- | Hash algorithms supported in certificates.

--

-- This relates the typed hash algorithm @hash@ to the 'HashALG' value.

data GHash hash = GHash { getHashALG :: HashALG, getHashAlgorithm :: hash }

hashMD2    :: GHash MD2

hashMD5    :: GHash MD5

hashSHA1   :: GHash SHA1

hashSHA224 :: GHash SHA224

hashSHA256 :: GHash SHA256

hashSHA384 :: GHash SHA384

hashSHA512 :: GHash SHA512

hashMD2    = GHash HashMD2    MD2

hashMD5    = GHash HashMD5    MD5

hashSHA1   = GHash HashSHA1   SHA1

hashSHA224 = GHash HashSHA224 SHA224

hashSHA256 = GHash HashSHA256 SHA256

hashSHA384 = GHash HashSHA384 SHA384

hashSHA512 = GHash HashSHA512 SHA512

-- | Signature and hash algorithms instantiated with parameters.

data Alg pub priv where

    AlgRSA    :: (HashAlgorithm hash, RSA.HashAlgorithmASN1 hash)

              => Int

              -> GHash hash

              -> Alg RSA.PublicKey RSA.PrivateKey

    AlgRSAPSS :: HashAlgorithm hash

              => Int

              -> PSS.PSSParams hash B.ByteString B.ByteString

              -> GHash hash

              -> Alg RSA.PublicKey RSA.PrivateKey

    AlgDSA    :: HashAlgorithm hash

              => DSA.Params

              -> GHash hash

              -> Alg DSA.PublicKey DSA.PrivateKey

    AlgEC     :: HashAlgorithm hash

              => ECC.CurveName

              -> GHash hash

              -> Alg ECDSA.PublicKey ECDSA.PrivateKey

-- | Types of public and private keys used by a signature algorithm.

type Keys pub priv = (Alg pub priv, pub, priv)

-- | Generates random keys for a signature algorithm.

generateKeys :: Alg pub priv -> IO (Keys pub priv)

generateKeys alg@(AlgRSA bits      _) = generateRSAKeys alg bits

generateKeys alg@(AlgRSAPSS bits _ _) = generateRSAKeys alg bits

generateKeys alg@(AlgDSA params    _) = do

    x <- DSA.generatePrivate params

    let y = DSA.calculatePublic params x

    return (alg, DSA.PublicKey params y, DSA.PrivateKey params x)

generateKeys alg@(AlgEC name       _) = do

    let curve = ECC.getCurveByName name

    (pub, priv) <- ECC.generate curve

    return (alg, pub, priv)

generateRSAKeys :: Alg RSA.PublicKey RSA.PrivateKey

                -> Int

                -> IO (Alg RSA.PublicKey RSA.PrivateKey, RSA.PublicKey, RSA.PrivateKey)

generateRSAKeys alg bits = addAlg <$> RSA.generate size e

  where

    addAlg (pub, priv) = (alg, pub, priv)

    size = bits `div` 8

    e    = 3

getPubKey :: Alg pub priv -> pub -> PubKey

getPubKey (AlgRSA    _    _) key = PubKeyRSA key

getPubKey (AlgRSAPSS _ _  _) key = PubKeyRSA key

getPubKey (AlgDSA    _    _) key = PubKeyDSA key

getPubKey (AlgEC     name _) key = PubKeyEC (PubKeyEC_Named name pub)

  where

    ECC.Point x y = ECDSA.public_q key

    pub   = SerializedPoint bs

    bs    = B.cons 4 (i2ospOf_ bytes x `B.append` i2ospOf_ bytes y)

    bits  = ECC.curveSizeBits (ECC.getCurveByName name)

    bytes = (bits + 7) `div` 8

getSignatureALG :: Alg pub priv -> SignatureALG

getSignatureALG (AlgRSA    _   hash) = SignatureALG (getHashALG hash) PubKeyALG_RSA

getSignatureALG (AlgRSAPSS _ _ hash) = SignatureALG (getHashALG hash) PubKeyALG_RSAPSS

getSignatureALG (AlgDSA    _   hash) = SignatureALG (getHashALG hash) PubKeyALG_DSA

getSignatureALG (AlgEC     _   hash) = SignatureALG (getHashALG hash) PubKeyALG_EC

doSign :: Alg pub priv -> priv -> B.ByteString -> IO B.ByteString

doSign (AlgRSA _ hash)        key msg = do

    result <- RSA.signSafer (Just $ getHashAlgorithm hash) key msg

    case result of

        Left err      -> error ("doSign(AlgRSA): " ++ show err)

        Right sigBits -> return sigBits

doSign (AlgRSAPSS _ params _) key msg = do

    result <- PSS.signSafer params key msg

    case result of

        Left err      -> error ("doSign(AlgRSAPSS): " ++ show err)

        Right sigBits -> return sigBits

doSign (AlgDSA _ hash)        key msg = do

    sig <- DSA.sign key (getHashAlgorithm hash) msg

    return $ encodeASN1' DER

                 [ Start Sequence

                 , IntVal (DSA.sign_r sig)

                 , IntVal (DSA.sign_s sig)

                 , End Sequence

                 ]

doSign (AlgEC _ hash)         key msg = do

    sig <- ECDSA.sign key (getHashAlgorithm hash) msg

    return $ encodeASN1' DER

                 [ Start Sequence

                 , IntVal (ECDSA.sign_r sig)

                 , IntVal (ECDSA.sign_s sig)

                 , End Sequence

                 ]

-- Certificate utilities --

-- | Holds together a certificate and its private key for convenience.

--

-- Contains also the crypto algorithm that both are issued from.  This is

-- useful when signing another certificate.

data Pair pub priv = Pair

    { pairAlg        :: Alg pub priv

    , pairSignedCert :: SignedCertificate

    , pairKey        :: priv

    }

-- | Builds a DN with a single component.

mkDn :: String -> DistinguishedName

mkDn cn = DistinguishedName [(getObjectID DnCommonName, fromString cn)]

-- | Used to build a certificate extension.

mkExtension :: Extension a => Bool -> a -> ExtensionRaw

mkExtension crit ext = ExtensionRaw (extOID ext) crit (extEncodeBs ext)

-- | Default extensions in leaf certificates.

leafStdExts :: [ExtensionRaw]

leafStdExts = [ku, eku]

  where

    ku  = mkExtension False $ ExtKeyUsage

               [ KeyUsage_digitalSignature , KeyUsage_keyEncipherment ]

    eku = mkExtension False $ ExtExtendedKeyUsage

               [ KeyUsagePurpose_ServerAuth , KeyUsagePurpose_ClientAuth ]

-- Authority signing a certificate --

--

-- When the certificate is self-signed, issuer and subject are the same.  So

-- they have identical signature algorithms.  The purpose of the GADT is to

-- hold this constraint only in the self-signed case.

-- | Authority signing a certificate, itself or another certificate.

data Auth pubI privI pubS privS where

    Self :: (pubI ~ pubS, privI ~ privS) => Auth pubI privI pubS privS

    CA   ::              Pair pubI privI -> Auth pubI privI pubS privS

foldAuth :: a

         -> (Pair pubI privI -> a)

         -> Auth pubI privI pubS privS

         -> a

foldAuth x _ Self   = x          -- no constraint used

foldAuth _ f (CA p) = f p

foldAuthPriv :: privS

             -> (Pair pubI privI -> privI)

             -> Auth pubI privI pubS privS

             -> privI

foldAuthPriv x _ Self   = x      -- uses constraint privI ~ privS

foldAuthPriv _ f (CA p) = f p

foldAuthPubPriv :: k pubS privS

                -> (Pair pubI privI -> k pubI privI)

                -> Auth pubI privI pubS privS

                -> k pubI privI

foldAuthPubPriv x _ Self   = x   -- uses both constraints

foldAuthPubPriv _ f (CA p) = f p

-- Certificate creation functions --

-- | Builds a certificate using the supplied keys and signs it with an

-- authority (itself or another certificate).

mkCertificate :: Int                        -- ^ Certificate version

              -> Integer                    -- ^ Serial number

              -> DistinguishedName          -- ^ Subject DN

              -> (DateTime, DateTime)       -- ^ Certificate validity period

              -> [ExtensionRaw]             -- ^ Extensions to include

              -> Auth pubI privI pubS privS -- ^ Authority signing the new certificate

              -> Keys pubS privS            -- ^ Keys for the new certificate

              -> IO (Pair pubS privS)       -- ^ The new certificate/key pair

mkCertificate version serial dn validity exts auth (algS, pubKey, privKey) = do

    signedCert <- objectToSignedExactF signatureFunction cert

    return Pair { pairAlg        = algS

                , pairSignedCert = signedCert

                , pairKey        = privKey

                }

  where

    pairCert = signedObject . getSigned . pairSignedCert

    cert = Certificate

        { certVersion      = version

        , certSerial       = serial

        , certSignatureAlg = signAlgI

        , certIssuerDN     = issuerDN

        , certValidity     = validity

        , certSubjectDN    = dn

        , certPubKey       = getPubKey algS pubKey

        , certExtensions   = extensions

        }

    signingKey = foldAuthPriv     privKey pairKey auth

    algI       = foldAuthPubPriv  algS    pairAlg auth

    signAlgI   = getSignatureALG algI

    issuerDN   = foldAuth dn (certSubjectDN . pairCert) auth

    extensions = Extensions (if null exts then Nothing else Just exts)

    signatureFunction objRaw = do

        sigBits <- doSign algI signingKey objRaw

        return (sigBits, signAlgI)

-- | Builds a CA certificate using the supplied keys and signs it with an

-- authority (itself or another certificate).

mkCA :: Integer                    -- ^ Serial number

     -> String                     -- ^ Common name

     -> (DateTime, DateTime)       -- ^ CA validity period

     -> Maybe ExtBasicConstraints  -- ^ CA basic constraints

     -> Maybe ExtKeyUsage          -- ^ CA key usage

     -> Auth pubI privI pubS privS -- ^ Authority signing the new certificate

     -> Keys pubS privS            -- ^ Keys for the new certificate

     -> IO (Pair pubS privS)       -- ^ The new CA certificate/key pair

mkCA serial cn validity bc ku =

    let exts = catMaybes [ mkExtension True <$> bc, mkExtension False <$> ku ]

    in mkCertificate 2 serial (mkDn cn) validity exts

-- | Builds a leaf certificate using the supplied keys and signs it with an

-- authority (itself or another certificate).

mkLeaf :: String                     -- ^ Common name

       -> (DateTime, DateTime)       -- ^ Certificate validity period

       -> Auth pubI privI pubS privS -- ^ Authority signing the new certificate

       -> Keys pubS privS            -- ^ Keys for the new certificate

       -> IO (Pair pubS privS)       -- ^ The new leaf certificate/key pair

mkLeaf cn validity = mkCertificate 2 100 (mkDn cn) validity leafStdExts