{-# LANGUAGE OverloadedStrings #-}
module TLSA
    ( TLSA(..)
    , CertUsage(..)
    , Selector(..)
    , MatchingType(..)
    , fromByteString
    , toByteString
    , match
    ) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as L
import qualified Crypto.Hash.SHA256 as SHA256
import qualified Crypto.Hash.SHA512 as SHA512
import Data.X509          ( certPubKey, Certificate )
import Data.ASN1.Types    ( toASN1 )
import Data.ASN1.Encoding ( encodeASN1 )
import Data.ASN1.BinaryEncoding
import Control.Applicative
import Data.Word
import Data.Maybe
import Data.Monoid

{- INLINE fromWord8 #-}
fromWord8 :: Enum a => Word8 -> a
fromWord8 = toEnum . fromEnum

{- INLINE toWord8 -}
toWord8 :: Enum a => a -> Word8
toWord8 = toEnum . fromEnum

data CertUsage
    = CAConstraint
    | ServiceCertificateConstraint
    | TrustAnchorAssertion
    | DomainIssued
    | CertUsage Word8
 deriving (Eq, Ord, Show, Read)

instance Enum CertUsage where
    fromEnum CAConstraint                 = 0
    fromEnum ServiceCertificateConstraint = 1
    fromEnum TrustAnchorAssertion         = 2
    fromEnum DomainIssued                 = 3
    fromEnum (CertUsage n)                = fromEnum n
    toEnum 0 = CAConstraint
    toEnum 1 = ServiceCertificateConstraint
    toEnum 2 = TrustAnchorAssertion
    toEnum 3 = DomainIssued
    toEnum n = CertUsage (toEnum n)

data Selector
    = FullCertificate
    | SubjectPublicKeyInfo
    | Selector Word8  -- Selector 255 reserved for private use
 deriving (Eq,Ord,Show)

instance Enum Selector where
    fromEnum FullCertificate      = 0
    fromEnum SubjectPublicKeyInfo = 1
    fromEnum (Selector n)         = fromEnum n
    toEnum 0   = FullCertificate
    toEnum 1   = SubjectPublicKeyInfo
    toEnum n = Selector (toEnum n)

data MatchingType
    = Match_Exact
    | Match_SHA256
    | Match_SHA512
    | Match Word8 -- Match 255 reserved for private use
 deriving (Eq,Ord,Show)

instance Enum MatchingType where
    fromEnum Match_Exact  = 0
    fromEnum Match_SHA256 = 1
    fromEnum Match_SHA512 = 2
    fromEnum (Match n)    = fromEnum n
    toEnum 0 = Match_Exact
    toEnum 1 = Match_SHA256
    toEnum 2 = Match_SHA512
    toEnum n = Match (toEnum n)


data TLSA = TLSA
    { tlsaCertUsage       :: CertUsage
    , tlsaSelector        :: Selector
    , tlsaMatchingType    :: MatchingType
    , tlsaAssociationData :: BS.ByteString
    }
 deriving (Eq, Ord, Show)

fromByteString :: BS.ByteString -> TLSA
fromByteString bs = TLSA (fromWord8 cu)
                         (fromWord8 sel)
                         (fromWord8 mat)
                         associationData
 where
    (csm,associationData) = BS.splitAt 3 bs
    (cu,sel,mat) =
        case BS.unpack csm of
            [cu,sel,mat] -> (cu,sel,mat)
            [cu,sel]     -> (cu,sel,0)
            [cu]         -> (cu,0,0)
            []           -> (0,0,0)

toByteString :: TLSA -> BS.ByteString
toByteString (TLSA cu sel mat associationData) = csm <> associationData
 where
    csm = BS.pack [ toWord8 cu
                  , toWord8 sel
                  , toWord8 mat ]

match :: Certificate -> TLSA -> Bool
match cert tlsa = fromMaybe False $
    (== tlsaAssociationData tlsa) <$> (hash <*> material)

 where
    encode obj = encodeASN1 DER (toASN1 obj [])

    material =
        case tlsaSelector tlsa of
            FullCertificate      -> Just $ encode cert
            SubjectPublicKeyInfo -> Just $ encode $ certPubKey cert
            _                    -> Nothing

    hash =
        case tlsaMatchingType tlsa of
            Match_Exact  -> Just L.toStrict
            Match_SHA256 -> Just SHA256.hashlazy
            Match_SHA512 -> Just SHA512.hashlazy
            _            -> Nothing