{-# LANGUAGE OverloadedStrings, ViewPatterns #-}
module TLSA
    ( TLSA(..)
    , CertUsage(..)
    , Selector(..)
    , MatchingType(..)
    , fromByteString
    , toByteString
    , match
    , certOrKey
    , IssuanceTest(..)
    , validate
    ) 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.ASN1.BinaryEncoding
import Control.Applicative
import Data.Word
import Data.Maybe
import Data.List
import Data.Monoid
import Data.Array.IArray

{- 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 is usage value 0 in RFC 6698.  Any CA certificate that 'match'es
    -- the 'TSLA' record is acceptable.  If the chain cannot be validated via
    -- 'TLSA' records alone (using 'TrustAnchorAssertion' or 'DomainIssued'),
    -- and there are any usable 'TLSA' records at all, then the 'validate'
    -- function will require at least one acceptable certificate.
    = AcceptableCA

    -- | This is usage value 1 in RFC 6698.  If the target certificate
    -- 'match'es, then it is considered acceptable. If the chain cannot be
    -- validated via 'TLSA' records alone (using 'TrustAnchorAssertion' or
    -- 'DomainIssued'), and there are any usable 'TLSA' records at all, then
    -- the 'validate' function will require at least one acceptable
    -- certificate.
    | AcceptableEE

    -- | 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
    -- 'AcceptableEE', PKIX validation should not be performed.
    | DomainIssued

    | CertUsage Word8 -- ^ Unusable.
 deriving (Eq, Ord, Show, Read)

instance Enum CertUsage where
    fromEnum AcceptableCA                 = 0
    fromEnum AcceptableEE = 1
    fromEnum TrustAnchorAssertion         = 2
    fromEnum DomainIssued                 = 3
    fromEnum (CertUsage n)                = fromEnum n
    toEnum 0 = AcceptableCA
    toEnum 1 = AcceptableEE
    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


-- comb [1,2,3] = [([],[1,2,3]),([1],[2,3]),([2,1],[3])]
comb :: [x] -> [([x],[x])]
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 rel x xs
--
-- Given a relation @rel@, a starting element @x@ and a collection of similarly
-- typed elements @xs@, returns the set of all ordered subsets @ks@ of @x:xs@
-- such that:
--
--  * @ x == 'head' ks @
--
--  * @ 'True' == 'and' \$ 'zipWith' rel ks ('tail' ks) @
-- 
-- The second condition is requiring consecutive pairs to satsify the given
-- relation @rel@.
--
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 (bs++as)

-- | These functions are used by 'validate' to determine when a certificate is
-- validly issued by another.
data IssuanceTest = IssuanceTest 
    { isIssuedBy :: SignedCertificate -> SignedCertificate -> Bool
    -- ^ This is used to validate a single link in a certificate chain.
    , isSignedBy :: SignedCertificate -> PubKey -> Bool
    -- ^ This is used to check signatures for trust anchor keys that are
    -- obtained from a 'TLSA' using 'certOrKey' but are not otherwise present
    -- in the input.
    }

-- | Use the the given set of 'TLSA' records to validate or paritally validate
-- a certificate given a list of other probably relevent certificates. Results
-- are interpreted as follows:
--
--   [@ Nothing   @] The certificate PASSED validation.
--
--   [@ Just \[\] @] The certificate FAILED validation.
--
--   [@ Just xss  @] A set of certificate issued-by chains.  If you trust any
--                   certificate in any of these chains, you may consider the
--                   certificate validated.  Otherwise, it failed validation.  
--
validate :: IssuanceTest -> [TLSA] -> SignedCertificate -> [SignedCertificate]
         -> Maybe [[SignedCertificate]]
validate (IssuanceTest isIssuedBy isSignedBy) rs cert chain
    | domainIssued         = Nothing
    | any hasAnchor chains = Nothing
    | null rs              = Just $ (certv !) .: chains
    | otherwise = Just $ (certv !) .: filter satisfiesConstraints chains
 where
    domainIssued = any (`match` cert) daneEEs

    threshold = length chain

    len = threshold + length anchor_certs

    certv :: Array Int SignedCertificate
    certv = listArray (0,len) $ cert:chain ++ anchor_certs

    (.:) = fmap . fmap

    a .<+ b  = (certv ! a)  `isIssuedBy` (certv ! b)

    a .<- b = (certv ! a) `isSignedBy` b

    isAnchor n = or [ n > threshold
                    , n `elem` anchors
                    , any (n .<-) anchor_keys ]

    chains = allChains (.<+) 0 [1..len]

    hasAnchor = any isAnchor

    ( anchor_certs, anchor_keys, anchors )
         = ( mapMaybe rightToMaybe absent
           , mapMaybe leftToMaybe  absent
           , fmap fst $ ns >>= snd )
      where
        (bs,ns) = partition (null . snd)
                    $ pairings (\r (_,c) -> match r c) daneTAs
                    $ zip [0..] (cert:chain)
        absent  = mapMaybe (certOrKey . fst) bs
        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` (certv ! c))
        caSatisfied r = any ((r `match`) . (certv !)) cs

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

pairings :: (a -> b -> Bool) -> [a] -> [b] -> [(a,[b])]
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)

-- | If a 'SignedCertificate' or a 'PubKey' is embedded in the 'TLSA' record
-- (i.e. 'matchingType' = 'Match_Exact'), then extract it.
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
                       }
    unsigned :: (Show a, ASN1Object a, Eq a) => a -> SignedExact a
    unsigned obj = fst $ objectToSignedExact fakeSign obj
     where fakeSign = const $ ("", SignatureALG_Unknown [], ())
-}