{-# LANGUAGE OverloadedStrings, ViewPatterns #-}
module TLSA
    ( TLSA(..)
    , CertUsage(..)
    , Selector(..)
    , MatchingType(..)
    , fromByteString
    , toByteString
    , match
    , Validation(..)
    , validate
    , validate2
    ) 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(..), SignedCertificate, getCertificate
                          , getSigned, getSignedData, signedObject, objectToSignedExact
                          , SignatureALG(SignatureALG_Unknown), SignedExact
                          , encodeSignedObject, decodeSignedObject ) -}
import Data.ASN1.Types    ( toASN1, ASN1Object, fromASN1 )
import Data.ASN1.Encoding ( encodeASN1, decodeASN1 )
import Data.Time.Clock.POSIX ( posixSecondsToUTCTime )

import Data.ASN1.BinaryEncoding
import Control.Applicative
import Data.Word
import Data.Maybe
import Data.List
import Data.Monoid

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

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

-- | The Certificate Usage Field as described in RFC 6698, section 2.1.1.
--
-- It is used by the 'validate' function in making a 'Validation' decision.
data CertUsage

    -- | This usage limits which CA can be used to issue certificates for a
    -- given service on a host.  PKIX-validated TLS connections for the domain
    -- should be considered invalid if the certification path does not include
    -- at least one certificate that 'match'es the 'TLSA' record.
    = CAConstraint

    -- | This usage limits which end entity certificate can be used by a given
    -- service on a host.  The TLS connection for the domain should be
    -- considered invalid if the end entity certificate does not 'match' the
    -- 'TLSA' record.
    | ServiceCertificateConstraint

    -- | This usage allows a domain name administrator to specify a new trust
    -- anchor.  This is useful if the domain issues its own certificates under
    -- its own CA that is not expected to be in the end users' collection of
    -- trust anchors.  When conducting PKIX validation for the domain, any
    -- certificate 'match'ing the 'TLSA' record can be treated as a trust
    -- anchor that does not require further validation.
    | TrustAnchorAssertion

    -- | This usage allows for a domain name administrator to issue
    -- certificates for a domain without involving a third-party CA.  The end
    -- entity certificate MUST 'match' the 'TLSA' record.  Unlike for a
    -- 'ServiceCertificateConstraint', PKIX validation should not be performed.
    | 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)

-- | Indicates what sort of object should be compared with 'associationData'.
data Selector
    = FullCertificate      -- ^ x.509 certificate
    | SubjectPublicKeyInfo -- ^ PKCS #8 formatted public key
    | Selector Word8       -- ^ value 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)

-- | Is 'associationData' an object of the form specified by 'Selector' or is
-- it only a hash of it?
data MatchingType
    = Match_Exact
    | Match_SHA256
    | Match_SHA512
    | Match Word8 -- ^ value 255 is 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)


-- | The parsed RDATA field of a TLSA DNS resource record (type 52) as
-- described in RFC 6698.
--
-- The 'match' function uses 'selector', 'matchingType' and 'associationData'
-- to implement a predicate on certificates obtained via the TLS protocol.  The
-- 'certUsage' field indicates what that predicate means.
data TLSA = TLSA
    { certUsage       :: CertUsage
    , selector        :: Selector
    , matchingType    :: MatchingType
    , associationData :: BS.ByteString
    }
 deriving (Eq, Ord, Show)

-- | Parse RDATA for a TLSA resource record.
fromByteString :: BS.ByteString -> TLSA
fromByteString bs = TLSA (fromWord8 cu)
                         (fromWord8 sel)
                         (fromWord8 mat)
                         dta
 where
    (csm,dta) = 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)

-- | Encode a valid RDATA field for a TLSA DNS record.
toByteString :: TLSA -> BS.ByteString
toByteString (TLSA cu sel mat dta) = csm <> dta
 where
    csm = BS.pack [ toWord8 cu
                  , toWord8 sel
                  , toWord8 mat ]

-- | Returns 'True' if the given certificate matches the given 'TLSA' object.
-- The algorithm for matching depends on the values of 'selector' and
-- 'matchingType' as described in RFC 6698.
match :: TLSA -> SignedCertificate -> Bool
match tlsa cert = fromMaybe False $
    (== associationData tlsa) <$> (hash <*> material)

 where
    key = encodeASN1 DER $ toASN1 keyobj []
        where keyobj = certPubKey $ getCertificate cert

    encoded_cert = L.fromChunks [encodeSignedObject cert]

    material :: Maybe L.ByteString
    material =
        case selector tlsa of
            FullCertificate      -> Just encoded_cert
            SubjectPublicKeyInfo -> Just key
            _                    -> Nothing

    hash :: Maybe (L.ByteString -> BS.ByteString)
    hash =
        case matchingType tlsa of
            Match_Exact  -> Just L.toStrict
            Match_SHA256 -> Just SHA256.hashlazy
            Match_SHA512 -> Just SHA512.hashlazy
            _            -> Nothing

unsigned :: (Show a, ASN1Object a, Eq a) => a -> SignedExact a
unsigned obj = fst $ objectToSignedExact fakeSign obj
 where fakeSign = const $ ("", SignatureALG_Unknown [], ())

data Validation
    = Failed [TLSA]
    -- ^ All of the given constraints failed.
    | TrustAnchors [SignedCertificate] [SignedCertificate]
    -- ^ Perform PKI validation with the given additional trust anchors.  The
    -- first list are trust anchors that occured in the chain.  The second list
    -- are anchors that did not occur in the chain.  If 'certVersion' == (-1)
    -- for any certificate in the second list, then all fields except
    -- 'certPubKey' should be ignored.
    | Passed
    -- ^ Valid end entity. Do not perform PKI.


-- comb [1,2,3] = [([],[1,2,3]),([1],[2,3]),([2,1],[3])]
comb cs = matchLength cs $ iterate sweepLeft ([],cs)
    where sweepLeft (xs,y:ys) = (y:xs,ys)
          matchLength = zipWith (flip const)

-- O(n²) worst case
-- O(n) best case
-- op is an antisymmetric operation
connectedChain op x [] = [x]
connectedChain op x xs = 
    case filter ((x `op`) . snd) $ comb xs of
        []          -> [x]
        (as,y:bs):_ -> x : connectedChain op y (bs++as)

allChains :: (a -> a -> Bool) -> a -> [a] -> [[a]]
allChains op x [] = [[x]]
allChains op x xs | null ts   = [[x]]
                  | otherwise = ts >>= f
 where
    ts = filter  ((x `op`) . head . snd) $ comb xs
    f (as,y:bs) = map (x:) $ allChains op y (as++bs)

-- | > validate2 isIssuedBy isSignedBy rs cert otherCerts
--
--   * isIssuedBy - validates one link in a certificate chain
--
--   * isSignedBy - True if the certificate is signed by the key
--
--   * rs - A set of TLSA records obtained from DNS
--
--   * cert - The end entity target cert to verify
--
--   * otherCerts - Other relevent certificates
--
--  Returns:
--
--   * @ Nothing @ - The end entity certificate is valid.
--
--   * @ Just chains @ - A set of connected certificate chains.
--                   The certificate is valid if and only if there
--                   exists a chain that contains a CA certificate
--                   that is issued by a trust anchor.
--
validate2 :: (SignedCertificate -> SignedCertificate -> Bool)
            -> (SignedCertificate -> PubKey -> Bool)
            -> [TLSA] -> SignedCertificate -> [SignedCertificate] -> Maybe [[SignedCertificate]]
validate2 isIssuedBy isSignedBy rs cert chain
    | some domainIssued    = Nothing
    | any hasAnchor chains = Nothing
    | otherwise            = Just $ filter satisfiesConstraints chains
  where
    domainIssued = filter (`match` cert) daneEEs

    chains = allChains isIssuedBy cert chain

    hasAnchor cs = some anchors
                    || some [ c | c <- cs, cert <- certs, c `isIssuedBy` cert ]
                    || some [ c | c <- cs, key <- keys, c `isSignedBy` key ]
     where
        (bs,ns) = partition (null . snd) $ pairings match daneTAs cs
        anchors = concatMap snd ns
        absent  = mapMaybe (certOrKey . fst) bs
        certs = mapMaybe rightToMaybe absent
        keys  = mapMaybe leftToMaybe  absent
        rightToMaybe (Right x) = Just x
        rightToMaybe _         = Nothing
        leftToMaybe  (Left x)  = Just x
        leftToMaybe  _         = Nothing

    satisfiesConstraints (c:cs) = any eeSatisfied pkixEEs
                                  || any caSatisfied pkixTAs
     where
        eeSatisfied = (`match` c)
        caSatisfied r = any (r `match`) cs

    some = not . null

    r .~ u = certUsage r == u
    (daneEEs,rs2) = partition (.~ DomainIssued)                 rs
    (daneTAs,rs3) = partition (.~ TrustAnchorAssertion)         rs2
    (pkixEEs,rs4) = partition (.~ ServiceCertificateConstraint) rs3
    (pkixTAs,_)   = partition (.~ CAConstraint)                 rs4

pairings op = loop
     where
        loop []     cs = []
        loop (m:ms) cs =
            case filter (op m . head . snd) $ comb cs of
                []          -> (m,[]):loop ms cs
                (as,b:bs):_ -> (m,[b]):loop ms (as++bs)

validate :: [TLSA] -> SignedCertificate -> [SignedCertificate] -> Validation
validate rs cert chain
    | some domainIssued               = Passed
    | null constraints || some passed = TrustAnchors anchors absent
    | otherwise                       = Failed constraints
  where
    domainIssued = filter (`match` cert) daneEEs
    (bs,ns) = partition (null . snd) $ pairings match daneTAs (cert:chain)
    anchors = concatMap snd ns
    absent  = mapMaybe (extractCert . fst) bs

    constraints = pkixEEs   ++ pkixTAs
    passed      = passedEEs ++ passedTAs
      where
        passedEEs = filter (`match` cert) pkixEEs
        -- TODO
        -- These passedTAs are only truly passed if the
        -- certs that match them are reachable.
        -- Where a cert is "reachable" if it is the end entity
        -- cert or if it is the issuer of a reachable cert.
        passedTAs = filter (`matchAny` (cert:chain)) pkixTAs
        matchAny t = any (t `match`)

    some = not . null

    r .~ u = certUsage r == u
    (daneEEs,rs2) = partition (.~ DomainIssued)                 rs
    (daneTAs,rs3) = partition (.~ TrustAnchorAssertion)         rs2
    (pkixEEs,rs4) = partition (.~ ServiceCertificateConstraint) rs3
    (pkixTAs,_)   = partition (.~ CAConstraint)                 rs4

certOrKey :: TLSA -> Maybe (Either PubKey SignedCertificate)
certOrKey tlsa@(matchingType->Match_Exact) =
    case selector tlsa of
        FullCertificate -> either (const Nothing)
                                  (Just . Right)
                                  (decodeSignedObject $ associationData tlsa)
        SubjectPublicKeyInfo -> do
            a <- hush $ decodeASN1 DER $ L.fromChunks [associationData tlsa]
            (key,_) <- hush $ fromASN1 a 
            return $ Left key
 where
        hush (Left _) = Nothing
        hush (Right a) = Just a

extractCert :: TLSA -> Maybe SignedCertificate
extractCert tlsa@(matchingType->Match_Exact) =
    case selector tlsa of
        FullCertificate -> either (const Nothing)
                                  Just
                                  (decodeSignedObject $ associationData tlsa)
        SubjectPublicKeyInfo -> do
            a <- hush $ decodeASN1 DER $ L.fromChunks [associationData tlsa]
            (key,_) <- hush $ fromASN1 a 
            return $ certificateFromKey key
 where
        hush (Left _) = Nothing
        hush (Right a) = Just a
extractCert _ = Nothing

certificateFromKey :: PubKey -> SignedCertificate
certificateFromKey key = unsigned cert
 where
    cert = Certificate { certPubKey = key
                       , certVersion = (-1)
                       , certSerial = 0
                       , certSignatureAlg = SignatureALG_Unknown []
                       , certIssuerDN = DistinguishedName []
                       , certValidity = ( posixSecondsToUTCTime (-1/0)
                                        , posixSecondsToUTCTime (1/0))
                       , certSubjectDN = DistinguishedName []
                       , certExtensions = Extensions Nothing
                       }