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

-- | Utilities for parsing a DMARC aggregated report, following RFC 7489. These functions also handle decompression (gzip or pkzip)

module Data.Mail.DMARC.Reports (
  -- * Parsers
  dmarcFeedbackFromStrict,
  dmarcFeedbackFromLazy,
  dmarcFeedbackFromStream,
  dmarcFeedbackFromXml,
  -- * Data types
  Feedback(..),
  ReportMetadata(..),
  Record(..),
  Row(..),
  PolicyPublished(..),
  PolicyEvaluated(..),
  DkimAuth(..),
  SpfAuth(..),
  AuthResult(..),
  IpAddress,
  AlignmentMode(..),
  PolicyAction(..),
  -- * Utility
  uncompressStream
  ) where

import Streaming
import Streaming.Zip
import qualified Streaming.Prelude as S
import qualified Streaming.ByteString as Q
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString as B
import qualified Codec.Archive.Zip as Z
import Control.Monad.IO.Class
import Control.Monad.Trans
import Control.Monad.Trans.Resource
import Text.XML.Light
import Text.Read (readMaybe)
import Data.Time.Clock
import Data.Time.Format
import Data.Maybe
import qualified Data.Text as T
import System.IO.Unsafe

-- | Uncompress gzip or pkzip stream. If it is neither gzipped or pkzipped, return the stream unmodified.
uncompressStream :: (MonadFail m, MonadIO m) => Q.ByteStream m r -> Q.ByteStream m ()
uncompressStream stream = do
  magic <- lift $ Q.toStrict_ $ Q.take 2 stream
  if magic == "PK"
    then pkunzip stream
    else if magic == "\x1F\x8b"
         then void $ gunzip stream
         else void stream
  where pkunzip stream = do
          lbs <- lift $ Q.toLazy_ stream
          let a = Z.toArchive lbs
          case Z.zEntries a of
            [] -> lift $ fail "Empty PKZip archive :("
            (e:_) -> Q.fromLazy $ Z.fromEntry e

data ReportMetadata = ReportMetadata {
  rmOrgName :: Maybe T.Text,
  rmEmail :: Maybe T.Text,
  rmExtraContactInfo :: Maybe T.Text,
  rmReportId :: Maybe Integer,
  rmDateRange :: Maybe (UTCTime, UTCTime)
  } deriving (Eq,Show,Read)

data AlignmentMode = Strict | Relaxed deriving (Eq)

instance Show AlignmentMode where
  show Strict = "s"
  show Relaxed = "r"

instance Read AlignmentMode where
  readsPrec _ r =
    [(Strict,s) | ("s",s) <- lex r]
    ++[(Relaxed,s) | ("r",s) <- lex r]

data PolicyAction = NoPolicy
                  | Quarantine
                  | Reject deriving (Eq)

instance Show PolicyAction where
  show NoPolicy = "none"
  show Quarantine = "quarantine"
  show Reject = "reject"

instance Read PolicyAction where
  readsPrec _ r =
    [(NoPolicy,s) | ("none",s) <- lex r]
    ++[(Quarantine,s) | ("quarantine",s) <- lex r]
    ++[(Reject,s) | ("reject",s) <- lex r]
  
data PolicyPublished = PolicyPublished {
  ppDomain :: T.Text,
  ppDkimAlignment :: AlignmentMode,
  ppSpfAlignment :: AlignmentMode,
  ppPolicy :: PolicyAction,
  ppSubdomainPolicy :: PolicyAction,
  ppPercent :: Int
  } deriving (Eq,Show,Read)

data AuthResult = NoAuth    -- ^ no policy published
                | AuthPass  -- ^ successful authentication
                | AuthFail  -- ^ authentication failed
                | TempError -- ^ temporary error in resolving the policy
                | PermError -- ^ permanent error in resolving the policy
                | AuthNeutral -- ^ ??? (DKIM only)
                | AuthPolicy  -- ^ ??? (DKIM only)
                | SoftFail  -- ^ auth failed, but won't reject (SPF only)
                deriving (Eq)

instance Show AuthResult where
  show NoAuth = "none"
  show AuthPass = "pass"
  show AuthFail = "fail"
  show TempError = "temperror"
  show PermError = "permerror"
  show AuthNeutral = "neutral"
  show AuthPolicy = "policy"
  show SoftFail = "softfail"

instance Read AuthResult where
  readsPrec _ r =
    [(NoAuth,s) | ("none",s) <- lex r]
    ++[(AuthPass,s) | ("pass",s) <- lex r]
    ++[(AuthFail,s) | ("fail",s) <- lex r]
    ++[(TempError,s) | ("temperror",s) <- lex r]
    ++[(PermError,s) | ("permerror",s) <- lex r]
    ++[(AuthNeutral,s) | ("neutral",s) <- lex r]
    ++[(AuthPolicy,s) | ("policy",s) <- lex r]
    ++[(SoftFail,s) | ("softfail",s) <- lex r]

type IpAddress = String

data PolicyEvaluated = PolicyEvaluated {
  peDisposition :: PolicyAction,
  peDkimResult :: AuthResult,
  peSpfResult :: AuthResult
  } deriving (Eq, Show, Read)

data Row = Row {
  rwSourceIp :: IpAddress,
  rwCount :: Int,
  rwPolicyEvaluated :: PolicyEvaluated
  } deriving (Eq, Show, Read)

data Record = Record {
  rcRow :: Row,
  rcAuthDkim :: Maybe DkimAuth,
  rcAuthSpf :: Maybe SpfAuth
  } deriving (Eq, Show, Read)

data DkimAuth = DkimAuth {
  daResult :: AuthResult,
  daDomain :: T.Text,
  daSelector :: T.Text
  } deriving (Eq, Show, Read)

data SpfAuth = SpfAuth {
  saResult :: AuthResult,
  saDomain :: T.Text
  } deriving (Eq, Show, Read)
  
data Feedback = Feedback {
  fbReportMetadata :: ReportMetadata,
  fbPolicyPublished :: PolicyPublished,
  fbRecords :: [Record]
  } deriving (Eq,Show,Read)

-- | Gets an XML document root from a streaming bytestring
xmlFromStream :: (MonadFail m, MonadIO m) => Q.ByteStream m r -> m (Maybe Element)
xmlFromStream stream = do
  let stream' = uncompressStream stream
  lbs <- Q.toLazy_ stream'
  return $ parseXMLDoc lbs

test :: IO ()
test = do
  let filename = "../google.com!enumeration.eu!1607299200!1607385599.zip"
  runResourceT $ do
    mfb <- dmarcFeedbackFromStream $ Q.readFile filename
    case mfb of
      Nothing -> liftIO $ putStrLn "nothing :("
      Just x -> liftIO $ putStrLn $ show x

-- | Parse a DMARC report from a 'Q.ByteStream'
dmarcFeedbackFromStream :: (MonadFail m, MonadIO m) => Q.ByteStream m r -> m (Maybe Feedback)
dmarcFeedbackFromStream stream = do
  mxml <- xmlFromStream stream
  return $ mxml >>= processXml

-- | Parse a DMARC report from a strict 'B.ByteString'
dmarcFeedbackFromStrict :: B.ByteString -> Maybe Feedback
dmarcFeedbackFromStrict = unsafePerformIO . dmarcFeedbackFromStream . Q.fromStrict

-- | Parse a DMARC report from a lazy 'L.ByteString'
dmarcFeedbackFromLazy :: L.ByteString -> Maybe Feedback
dmarcFeedbackFromLazy = unsafePerformIO . dmarcFeedbackFromStream . Q.fromLazy

-- | Parse a DMARC report from an XML document
dmarcFeedbackFromXml = processXml

-- | Parse a DMARC report from an XML document (also: 'dmarcFeedbackFromXml')
processXml :: Element -> Maybe Feedback
processXml root =
  let strContent' = T.pack . strContent
      reportMetadata = findChild (unqual "report_metadata") root
      policyPublished = findChild (unqual "policy_published") root
      rmOrgName = strContent' <$> (findChild (unqual "org_name") =<< reportMetadata)
      rmEmail = strContent' <$> (findChild (unqual "email") =<< reportMetadata)
      rmExtraContactInfo = strContent' <$> (findChild (unqual "extra_contact_info") =<< reportMetadata)
      rmReportId = readMaybe =<< strContent <$> (findChild (unqual "report_id") =<< reportMetadata)
      rmDateRange = do
        rm <- reportMetadata
        rangee <- findChild (unqual "date_range") rm
        begine <- findChild (unqual "begin") rangee
        ende <- findChild (unqual "end") rangee
        let begins = strContent begine
            ends = strContent ende
        begint <- parseTimeM True defaultTimeLocale "%s" begins
        endt <- parseTimeM True defaultTimeLocale "%s" ends
        return (begint, endt)
      mpolicy = do
        pp <- policyPublished
        ppDomain <- strContent' <$> findChild (unqual "domain") pp
        ppDkimAlignment <- (do
          adkim <- findChild (unqual "adkim") pp
          readMaybe $ strContent adkim) <|> Just Relaxed
        ppSpfAlignment <- (do
          aspf <- findChild (unqual "spf") pp
          readMaybe $ strContent aspf) <|> Just Relaxed
        ppPolicy <- (do
          p <- findChild (unqual "p") pp
          readMaybe $ strContent p) <|> Just NoPolicy
        ppSubdomainPolicy <- (do
          sp <- findChild (unqual "sp") pp
          readMaybe $ strContent sp) <|> Just NoPolicy
        ppPercent <- (do
          pct <- findChild (unqual "pct") pp
          readMaybe $ strContent pct) <|> Just 100
        return PolicyPublished{..}
      mpe el = do
        dispo <- (readMaybe =<< strContent <$> findChild (unqual "disposition") el) <|> Just NoPolicy
        dkim <- (readMaybe =<< strContent <$> findChild (unqual "dkim") el) <|> Just NoAuth
        spf <- (readMaybe =<< strContent <$> findChild (unqual "spf") el) <|> Just NoAuth
        return $ PolicyEvaluated dispo dkim spf
      mrow el = do
        pe <- (mpe =<< findChild (unqual "policy_evaluated") el) <|> Just (PolicyEvaluated NoPolicy NoAuth NoAuth)
        sourceip <- strContent <$> findChild (unqual "source_ip") el
        count <- readMaybe =<< (strContent <$> findChild (unqual "count") el)
        return $ Row sourceip count pe
      mda el = do
        domain <- strContent' <$> findChild (unqual "domain") el
        result <- readMaybe =<< strContent <$> findChild (unqual "result") el
        selector <- strContent' <$> findChild (unqual "selector") el
        return $ DkimAuth result domain selector
      msa el = do
        domain <- strContent' <$> findChild (unqual "domain") el
        result <- readMaybe =<< strContent <$> findChild (unqual "result") el
        return $ SpfAuth result domain
      mrecord el = do
        row <- mrow =<< findChild (unqual "row") el
        let authresults = findChild (unqual "auth_results") el
        let da = mda =<< findChild (unqual "dkim") =<< authresults
        let sa = msa =<< findChild (unqual "spf") =<< authresults
        return $ Record row da sa
  in do
    policy <- mpolicy
    let records = mapMaybe mrecord $ findChildren (unqual "record") root
    return $ Feedback ReportMetadata{..} policy records