-- |
-- Module      : Network.DNS
-- Description : Generic DNS utilities
--
-- There is no standardized presentation and parsing format for domain names.
-- In this library we assume a domain name and pattern to be specified as a text with an ASCII dot @.@ acting as a separator and terminator.
-- We do not admit arbitrary unicode codepoints, only the following subset of ASCII is acceptable per label:
-- [a-z], [A-Z], [0-9], '_', '-', '*'
--
-- Punycoding, if desired, must be taken care of the user.
--
-- In addition, we allow a backslash to be used as an escaping character for the following possible sequences:
--
-- Escape sequences
-- The domain name and pattern language here allows for the following escape sequences
--
-- @
--    \\.      gives a dot inside a label, rather than a label separator
--    \\\\      gives a backslash inside a label
--    \\012    gives an arbitrary octet inside a label as specified by the three octets
-- @
--
-- For example: @foo\\.bar.quux.@ is a domain name comprised of two labels @foo.bar@ and @quux@
module Network.DNS
  ( Domain
  , getDomain
  , mkDomain
  , mkDomain'
  , DomainLabel
  , getDomainLabel
  , getDomainLabelCF
  , mkDomainLabel
  , unsafeMkDomainLabel
  , unsafeSingletonDomainLabel
  , foldCase
  , foldCase_
  , foldCaseLabel

  -- * Parsing
  , parseAbsDomain
  , parseAbsDomainRelax
  , parseDomainLabel
  , absDomainP
  , absDomainRelaxP
  , domainLabelP

  -- * Pretty printing
  , pprDomain_
  , pprDomain
  , pprDomainCF
  , pprDomainLabel
  , pprDomainLabelCF
  )
where

import           Data.Char (isDigit, isLower, isUpper)
import           Data.Coerce (coerce)
import           Data.Word (Word8)

import           Control.Applicative.Combinators
import           Control.Monad (when)
import           Data.ByteString.Internal (w2c)
import           Data.Char (ord)
import           Data.Foldable (asum)

import           Data.Attoparsec.Text ((<?>))
import qualified Data.Attoparsec.Text as A
import           Data.ByteString.Internal (c2w)
import qualified Data.ByteString.Short as BS
import qualified Data.Text as T
import qualified Data.Text.Encoding as T

import           Network.DNS.Internal

-- | Parse an absolute domain. Convenience wrapper for 'absDomainP'.
parseAbsDomain :: T.Text -> Either String Domain
parseAbsDomain = A.parseOnly (absDomainP <* A.endOfInput)

-- | Parse a singular domain label. Convenience wrapper for 'domainLabelP'.
parseDomainLabel :: T.Text -> Either String DomainLabel
parseDomainLabel = A.parseOnly (domainLabelP <* A.endOfInput)

-- | Version of parseAbsDomain that also considers a domain name without a trailing dot
-- to be absolute.
parseAbsDomainRelax :: T.Text -> Either String Domain
parseAbsDomainRelax = A.parseOnly (absDomainRelaxP <* A.endOfInput)

-- | Turn a 'Domain' into a list of its labels.
--
-- prop> getDomain . mkDomain ~~~ id
-- prop> mkDomain . getDomain ~~~ id
getDomain :: Domain -> [DomainLabel]
getDomain = coerce

-- | Turn a list of labels into a 'Domain'.
--
-- prop> getDomain . mkDomain ~~~ id
-- prop> mkDomain . getDomain ~~~ id
mkDomain :: [DomainLabel] -> Domain
mkDomain = coerce

-- | Turn a list of text labels into a 'Domain'
--
-- Codepoints outside ASCII are officially not supported.
mkDomain' :: [T.Text] -> Domain
mkDomain' xs = mkDomain (mkDomainLabel . BS.toShort . T.encodeUtf8 <$> xs)

-- | Get the wire-representation of a domain label.
{-# INLINE getDomainLabel #-}
getDomainLabel :: DomainLabel -> BS.ShortByteString
getDomainLabel = getDomainLabel_

-- | Get the [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3) case-folded wire-representation of a domain label.
{-# INLINE getDomainLabelCF #-}
getDomainLabelCF :: DomainLabel -> BS.ShortByteString
getDomainLabelCF = getDomainLabelCF_

-- | Smart constructor for 'DomainLabel'
mkDomainLabel :: BS.ShortByteString -> DomainLabel
mkDomainLabel l = DomainLabel l (sbsMap foldCase_ l)

-- | Unsafely construct a 'DomainLabel'. The argument must already be case-folded according to [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
unsafeMkDomainLabel :: BS.ShortByteString -> DomainLabel
unsafeMkDomainLabel l = DomainLabel l l


-- | Unsafely construct a 'DomainLabel' from a single Word8. The argument must already be case-folded according to [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
unsafeSingletonDomainLabel :: Word8 -> DomainLabel
unsafeSingletonDomainLabel l = DomainLabel (sbsSingleton l) (sbsSingleton l)

-- | Case-folding of a domain according to [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
-- Note 'Domain' will memoize a case-folded variant for 'Eq', 'Ord' and pretty printing already. This function is not useful to most.
foldCase :: Domain -> Domain
foldCase (Domain ls) = Domain (foldCaseLabel <$> ls)

-- | Case-folding of a domain label according to [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
-- Note 'DomainLabel' will memoize a case-folded variant for 'Eq', 'Ord' and pretty printing already. This function is not useful to most.
{-# INLINE foldCaseLabel #-}
foldCaseLabel :: DomainLabel -> DomainLabel
foldCaseLabel (DomainLabel _l cf) = DomainLabel cf cf

{-# INLINE foldCase_ #-}
foldCase_ :: Word8 -> Word8
foldCase_ w | c2w 'A' <= w && w <= c2w 'Z'
            = w + 0x20
            | otherwise
            = w

-- | Print an arbitrary domain into a presentation format.
--
-- This function nearly roundtrips with 'parseAbsDomain' up to escape sequence equivalence
--
-- prop> parseAbsDomain . pprDomain ~~~ id
pprDomain :: Domain -> T.Text
pprDomain (Domain l) = T.pack (fromDList build)
  where
    build :: DList Char
    build = foldr (\x buf -> buildLabel x <> singleton '.' <> buf) mempty l

pprDomain_ :: Domain -> T.Text
pprDomain_ (Domain ls) = pprLabelsUtf16 (getDomainLabel_ <$> ls)

-- | Print an arbitrary domain into a presentation format after case-folding according to [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
--
-- This function nearly roundtrips with 'parseAbsDomain' up to escape sequence equivalence and case folding.
--
-- prop> parseAbsDomain . pprDomainCF ~~~ id
pprDomainCF :: Domain -> T.Text
pprDomainCF (Domain l) = T.pack (fromDList build)
  where
    build :: DList Char
    build = foldr (\x buf -> buildLabelCF x <> singleton '.' <> buf) mempty l

-- | Print a singular domain label into a presentation format.
pprDomainLabel :: DomainLabel -> T.Text
pprDomainLabel = T.pack . fromDList . buildLabel

-- | Print a singular domain label into a presentation format after case-folding according to [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
pprDomainLabelCF :: DomainLabel -> T.Text
pprDomainLabelCF = T.pack . fromDList . buildLabelCF


-- | Attoparsec 'A.Parser' for absolute domains. See 'parseAbsDomainRelax' for a convenience warpper.
-- This variant differs from 'absDomainP' in that it does not care whether the domain
-- name ends in a dot.
absDomainRelaxP :: A.Parser Domain
absDomainRelaxP = do
    d <- go
    let l = encodedLength d
    when (l >= 255) (fail "domain name too long")
    pure d

  where
    go = Domain <$> domainLabelP `sepBy1` A.char '.' <* optional (A.char '.')

-- | Calculate the wire-encoded length of a domain name.
encodedLength :: Domain -> Int
encodedLength (Domain labels) = sum (BS.length <$> l') + length l'
  where
    l' :: [BS.ShortByteString]
    l' = getDomainLabel <$> labels

-- | Attoparsec 'A.Parser' for absolute domains. See 'parseAbsDomain' for a convenience wrapper.
-- For a parser that also admits domain forms without a leading dot, see 'absDomainRelaxP'.
absDomainP :: A.Parser Domain
absDomainP = do
    d <- go
    let l = encodedLength d
    when (l >= 255) (fail "domain name too long")
    pure d

  where
    go = Domain <$> asum
        [ domainLabelP `endBy1` A.char '.'
        , [] <$ A.char '.' -- The root domain itself
        ]

-- | Predicate selecting characters allowed in a domain label without escaping.
{-# INLINE isLabelChar #-}
isLabelChar :: Char -> Bool
isLabelChar x = isLower x || isDigit x || isUpper x || x == '-' || x == '_' || x == '*'

-- | Attoparsec 'A.Parser' for a singular domain label. See 'parseDomainLabel' for a convenince wrapper. Also see 'absDomainP'.
domainLabelP :: A.Parser DomainLabel
domainLabelP = mkDomainLabel . BS.pack <$> (some labelChar)
  where
    labelChar :: A.Parser Word8
    labelChar = do
        c <- A.satisfy (\x -> isLabelChar x || x == '\\') <?> "domain label character"
        case c of
            '\\' -> escape
            _    -> pure (c2w c)

    escape :: A.Parser Word8
    escape = asum [ c2w <$> A.char '.'
                  , c2w <$> A.char '\\'
                  , octal ]
             <?> "escapable character"

    octal :: A.Parser Word8
    octal = do
        o1 <- v <$> A.satisfy isOctal
        o2 <- v <$> A.satisfy isOctal
        o3 <- v <$> A.satisfy isOctal
        pure (fromIntegral (o1 * 64 +
                            o2 * 8 +
                            o3))

        where
            v c = ord c - 48

            isOctal :: Char -> Bool
            isOctal c = c >= '0' && c <= '7'

-- | Make a case-folded string from a 'DomainLabel' suitable for pretty printing
{-# INLINE buildLabelCF #-}
buildLabelCF :: DomainLabel -> DList Char
buildLabelCF = buildLabel_ . getDomainLabelCF_

-- | Make a string from a 'DomainLabel' suitable for pretty printing
{-# INLINE buildLabel #-}
buildLabel :: DomainLabel -> DList Char
buildLabel = buildLabel_ . getDomainLabel_

{-# INLINE buildLabel_ #-}
buildLabel_ :: BS.ShortByteString -> DList Char
buildLabel_ bs = toDList (replace (BS.unpack bs))
  where
    {-# INLINE replace #-}
    replace :: [Word8] -> [Char]
    replace (x:xs) = case x of
      _ | isLabelChar (w2c x) -> (w2c x) : replace xs

      0x2e -> '\\' : '.' : replace xs
      0x5c -> '\\' : '\\' : replace xs
      _    -> '\\' : o1 : o2 : o3 : replace xs
        where
            (o1, o2, o3) = case quotRem x 8 of
                (v1, r3) -> case quotRem v1 8 of
                    (v2, r2) -> case quotRem v2 8 of
                        (_, r1)  -> (showD r1, showD r2, showD r3)
    replace [] = []
    {-# INLINE showD #-}
    showD :: Word8 -> Char
    showD x = w2c (x + 0x30)