{-# LANGUAGE OverloadedStrings, FlexibleContexts, RecordWildCards #-}

-- | Utilities for parsing an SMTP-TLRPT aggregated report, following RFC 8460. These functions also handle decompression (gzip or pkzip)

module Data.Mail.TLSRPT.Reports (
  -- * Parsers
  tlsReportFromStrict,
  tlsReportFromLazy,
  tlsReportFromStream,
  tlsReportFromJson,
  -- * Data types
  Report(..),
  Policy(..),
  PolicyDesc(..),
  FailureDetails(..),
  IpAddress,
  PolicySummary(..),
  PolicyType(..),
  ResultType(..)
  ) where

import Data.Mail.DMARC.Reports (uncompressStream)
import Streaming
import qualified Streaming.Prelude as S
import Data.Aeson
import Data.Aeson.Types
import Data.Time.Clock
import Data.Time.Format
import qualified Data.Text as T
import Text.Read (readMaybe)
import Data.Aeson.Encoding (string)
import qualified Streaming.ByteString as Q
import Control.Monad.Trans.Resource
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import System.IO.Unsafe

type IpAddress = String

data PolicyType = NoPolicyFound -- ^ neither MTA-STS nor DANE
                | TlsaPolicy    -- ^ DANE policy (TLSA record)
                | StsPolicy     -- ^ MTA-STS policy
                deriving (Eq)

instance Show PolicyType where
  show NoPolicyFound = "no-policy-found"
  show TlsaPolicy = "tlsa"
  show StsPolicy = "sts"

instance Read PolicyType where
  readsPrec _ r =
    [(NoPolicyFound,s) | ("no-policy-found",s) <- lex' r]
    ++[(TlsaPolicy,s) | ("tlsa",s) <- lex' r]
    ++[(StsPolicy,s) | ("sts",s) <- lex' r]

instance ToJSON PolicyType where
  toJSON = toJSONShow
  toEncoding = toEncodingShow

toJSONShow :: Show a => a -> Value
toJSONShow = String . T.pack . show

toEncodingShow :: Show a => a -> Encoding
toEncodingShow = string . show

instance FromJSON PolicyType where
  parseJSON = parseJSONRead "PolicyType"

parseJSONRead :: Read a => String -> Value -> Parser a
parseJSONRead desc = withText desc $ \t ->
  case readMaybe (T.unpack t) of
    Nothing -> typeMismatch desc (String t)
    Just a -> return a

-- | Description of the evaluated policy
data PolicyDesc = PolicyDesc {
  -- | Policy type (DANE or MTA-STS)
  pdPolicyType :: PolicyType,
  -- | Applied policy as strings
  pdPolicyString :: [T.Text],
  -- | domain (if punycode: A-labels, not U-labels)
  pdPolicyDomain :: T.Text,
  -- | if MTA-STS, then this is a list of mx host patterns
  pdMxHost :: [T.Text]
  } deriving (Eq, Show, Read)

instance ToJSON PolicyDesc where
  toJSON PolicyDesc{..} =
    object ["policy-type" .= pdPolicyType,
            "policy-string" .= pdPolicyString,
            "policy-domain" .= pdPolicyDomain,
            "mx-host" .= pdMxHost]
  toEncoding PolicyDesc{..} =
    pairs ("policy-type" .= pdPolicyType <>
           "policy-string" .= pdPolicyString <>
           "policy-domain" .= pdPolicyDomain <>
           "mx-host" .= pdMxHost)

instance FromJSON PolicyDesc where
  parseJSON = withObject "PolicyDesc" $
    \o -> PolicyDesc
          <$> o .: "policy-type"
          <*> o .:? "policy-string" .!= []
          <*> o .: "policy-domain"
          <*> o .:? "mx-host" .!= []

-- | Summary of policy evaluation
data PolicySummary = PolicySummary {
  -- | Successful sessions
  psTotalSuccessfulSessionCount :: Int,
  -- | Failed sessions
  psTotalFailureSessionCount :: Int
  } deriving (Eq, Show, Read)

instance ToJSON PolicySummary where
  toJSON PolicySummary{..} =
    object ["total-successful-session-count" .= psTotalSuccessfulSessionCount,
            "total-failure-session-count" .= psTotalFailureSessionCount]
  toEncoding PolicySummary{..} =
    pairs ("total-successful-session-count" .= psTotalSuccessfulSessionCount <>
           "total-failure-session-count" .= psTotalFailureSessionCount)

instance FromJSON PolicySummary where
  parseJSON = withObject "PolicySummary" $
    \o -> PolicySummary
          <$> o .: "total-successful-session-count"
          <*> o .: "total-failure-session-count"

-- | Failure types
data ResultType = StarttlsNotSupported -- ^ recipient MX did not support STARTTLS
                | CertificateHostMismatch -- ^ MX hostname does not match any of the subject alternative names in the certificate
                | CertificateExpired -- ^ certificate has expired
                | CertificateNotTrusted -- ^ certificate is not signed by any trusted CA
                | ValidationFailure -- ^ general negotiation failure (other)
                | TlsaInvalid -- ^ validation error in the TLSA record
                | DnssecInvalid -- ^ no valid records were returned from the recursive resolver
                | DaneRequired -- ^ sending system is configured to require DANE TLSA records for all the MX hosts of the destination domain, but no DNSSEC-validated TLSA records were present for the MX host
                | StsPolicyFetchError -- ^ failure to retrieve an MTA-STS policy (e.g. host unreachable)
                | StsPolicyInvalid -- ^ validation error for the overall MTA-STS policy
                | StsWebpkiInvalid -- ^ MTA-STS policy could not be authenticated using PKIX validation
                deriving (Eq)

instance Show ResultType where
  show StarttlsNotSupported = "starttls-not-supported"
  show CertificateHostMismatch = "certificate-host-mismatch"
  show CertificateExpired = "certificate-expired"
  show CertificateNotTrusted = "certificate-not-trusted"
  show ValidationFailure = "validation-failure"
  show TlsaInvalid = "tlsa-invalid"
  show DnssecInvalid = "dnssec-invalid"
  show DaneRequired = "dane-required"
  show StsPolicyFetchError = "sts-policy-fetch-error"
  show StsPolicyInvalid = "sts-policy-invalid"
  show StsWebpkiInvalid = "sts-webpki-invalid"

instance Read ResultType where
  readsPrec _ r =
    [(rt,s)
    | (w,s) <- lex' r,
      let rt = case w of
            "starttls-not-supported" -> StarttlsNotSupported
            "certificate-host-mismatch" -> CertificateHostMismatch
            "certificate-expired" -> CertificateExpired
            "certificate-not-trusted" -> CertificateNotTrusted
            "validation-failure" -> ValidationFailure
            "tlsa-invalid" -> TlsaInvalid
            "dnssec-invalid" -> DnssecInvalid
            "dane-required" -> DaneRequired
            "sts-policy-fetch-error" -> StsPolicyFetchError
            "sts-policy-invalid" -> StsPolicyInvalid
            "sts-webpki-invalid" -> StsWebpkiInvalid]

lex' r = [span (\c -> c `elem` ('-':['a'..'z'])) r]
      
instance FromJSON ResultType where
  parseJSON = parseJSONRead "ResultType"

instance ToJSON ResultType where
  toJSON = toJSONShow
  toEncoding = toEncodingShow

data FailureDetails = FailureDetails {
  fdResultType :: ResultType,             -- ^ failure type
  fdSendingMtaIp :: IpAddress,            -- ^ address of the sending MTA
  fdReceivingMxHostname :: T.Text,        -- ^ hostname of the receiving MTA MX record
  fdReceivingMxHelo :: Maybe T.Text,      -- ^ the HELO or EHLO string from thte banner announced during the reported session
  fdReceivingIp :: IpAddress,             -- ^ destination IP address
  fdFailedSessionCount :: Int,            -- ^ number of attempted sessions that match the result type
  fdAdditionalInformation :: Maybe T.Text,-- ^ URI that points to additional information
  fdFailureReasonCode :: Maybe T.Text     -- ^ failure code, esp. useful if result type is validation-failure
  } deriving (Eq, Show, Read)

instance ToJSON FailureDetails where
  toJSON FailureDetails{..} =
    object (["result-type" .= fdResultType,
             "sending-mta-ip" .= fdSendingMtaIp,
             "receiving-mx-hostname" .= fdReceivingMxHostname,
             "receiving-ip" .= fdReceivingIp,
             "failed-session-count" .= fdFailedSessionCount]
             ++maybe [] (\h -> ["receiving-mx-helo" .= h]) fdReceivingMxHelo
             ++maybe [] (\h -> ["additional-information" .= h]) fdAdditionalInformation
             ++maybe [] (\h -> ["failure-reason-code" .= h]) fdFailureReasonCode)

instance FromJSON FailureDetails where
  parseJSON = withObject "FailureDetails" $
    \o -> FailureDetails
          <$> o .: "result-type"
          <*> o .: "sending-mta-ip"
          <*> o .: "receiving-mx-hostname"
          <*> o .:? "receiving-mx-helo"
          <*> o .: "receiving-ip"
          <*> o .: "failed-session-count"
          <*> o .:? "additional-information"
          <*> o .:? "failure-reason-code"

data Policy = Policy {
  -- | description of the evaluated policy
  pcPolicy :: PolicyDesc,
  -- | failure/success summary
  pcSummary :: PolicySummary,
  -- | details about failure types
  pcFailureDetails :: [FailureDetails]
  } deriving (Eq, Show, Read)

instance ToJSON Policy where
  toJSON Policy{..} =
    object ["policy" .= pcPolicy,
            "summary" .= pcSummary,
            "failure-details" .= pcFailureDetails]

instance FromJSON Policy where
  parseJSON = withObject "Policy" $
    \o -> Policy
          <$> o .: "policy"
          <*> o .: "summary"
          <*> o .:? "failure-details" .!= []

-- | The entire SMTP-TLSRPT report
data Report = Report {
  -- | organization responsible for the report
  rpOrganizationName :: T.Text,
  -- | start and end times for the report range
  rpDateRange :: (UTCTime, UTCTime),
  -- | e-mail address of the party responsible for the report
  rpContactInfo :: T.Text,
  -- | unique identifier
  rpReportId :: T.Text,
  -- | results for each policy
  rpPolicies :: [Policy]
  } deriving (Eq, Show, Read)

instance ToJSON Report where
  toJSON Report{..} =
    object ["organization-name" .= rpOrganizationName,
            "date-range" .= object
             ["start-datetime" .= fmt (fst rpDateRange),
              "end-datetime" .= fmt (snd rpDateRange)],
            "contact-info" .= rpContactInfo,
            "report-id" .= rpReportId,
            "policies" .= rpPolicies]
    where fmt = formatTime defaultTimeLocale "%FT%X%Z"
  
instance FromJSON Report where
  parseJSON = withObject "Report" $
    \o -> Report
          <$> o .: "organization-name"
          <*> mtime o
          <*> o .: "contact-info"
          <*> o .: "report-id"
          <*> o .: "policies"
    where mtime o = do
            range <- o .: "date-range"
            start <- range .: "start-datetime"
            end <- range .: "end-datetime"
            start' <- parseTimeM True defaultTimeLocale "%FT%X%Z" start
            end' <- parseTimeM True defaultTimeLocale "%FT%X%Z" end
            return (start', end')

-- | Parse a TLSRPT report from a JSON 'Value'
tlsReportFromJson :: Value -> Either String Report
tlsReportFromJson v =
  case fromJSON v of
    Error s -> Left s
    Success r -> Right r

-- | Parse a TLSRPT report from a strict 'B.ByteString'
tlsReportFromStrict :: B.ByteString -> Either String Report
tlsReportFromStrict = unsafePerformIO . tlsReportFromStream . Q.fromStrict

-- | Parse a TLSRPT report from a lazy 'L.ByteString'
tlsReportFromLazy :: L.ByteString -> Either String Report
tlsReportFromLazy = unsafePerformIO . tlsReportFromStream . Q.fromLazy

-- | Parse a TLSRPT report from a 'Q.ByteStream'
tlsReportFromStream :: (MonadFail m, MonadIO m) => Q.ByteStream m r -> m (Either String Report)
tlsReportFromStream stream = do
  let stream' = uncompressStream stream
  lbs <- Q.toLazy_ stream'
  return $ eitherDecode lbs

test :: IO ()
test = do
  let filename = "../google.com!enumeration.eu!1607990400!1608076799!001.json.gz"
  runResourceT $ do
    erpt <- tlsReportFromStream $ Q.readFile filename
    case erpt of
      Left e -> liftIO $ putStrLn e
      Right x -> liftIO $ putStrLn $ show x