{-# LANGUAGE OverloadedStrings, TransformListComp #-}

module RoundTripSpec where

import qualified Data.IP
import Data.IP (Addr, IP(..), IPv4, IPv6, toIPv4, toIPv6, makeAddrRange)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BS
import qualified Data.CaseInsensitive as CI
import Test.Hspec
import Test.Hspec.QuickCheck
import Test.QuickCheck (Gen, arbitrary, elements, forAll, frequency, listOf, oneof)
import GHC.Exts (the, groupWith)

import Network.DNS.Decode
import Network.DNS.Decode.Internal
import Network.DNS.Encode
import Network.DNS.Encode.Internal
import Network.DNS.Imports
import Network.DNS.Types

spec :: Spec
spec = do
    prop "IPv4" . forAll genIPv4 $ \ ip4 -> do
        let str = show ip4
        read str `shouldBe` ip4
        show (read str :: IPv4) `shouldBe` str

    prop "IPv6" . forAll genIPv6 $ \ ip6 -> do
        let str = show ip6
        read str `shouldBe` ip6
        show (read str :: IPv6) `shouldBe` str

    prop "TYPE" . forAll genTYPE $ \ t ->
        toTYPE (fromTYPE t) `shouldBe` t

    prop "Domain" . forAll genDomain $ \ dom -> do
        let bs = encodeDomain dom
        decodeDomain bs `shouldBe` Right dom
        fmap encodeDomain (decodeDomain bs) `shouldBe` Right bs

    prop "Mailbox" . forAll genMailbox $ \ dom -> do
        let bs = encodeMailbox dom
        decodeMailbox bs `shouldBe` Right dom
        fmap encodeMailbox (decodeMailbox bs) `shouldBe` Right bs

    prop "DNSFlags" . forAll (genDNSFlags 0x0f) $ \ flgs -> do
        let bs = encodeDNSFlags flgs
        decodeDNSFlags bs `shouldBe` Right flgs
        fmap encodeDNSFlags (decodeDNSFlags bs) `shouldBe` Right bs

    prop "ResourceRecord" . forAll genResourceRecord $ \ rr -> do
        let bs = encodeResourceRecord rr
        decodeResourceRecord bs `shouldBe` Right rr
        fmap encodeResourceRecord (decodeResourceRecord bs) `shouldBe` Right bs

    prop "DNSHeader" . forAll (genDNSHeader 0x0f) $ \ hdr ->
        decodeDNSHeader (encodeDNSHeader hdr) `shouldBe` Right hdr

    prop "DNSMessage" . forAll genDNSMessage $ \ msg ->
        decode (encode msg) `shouldBe` Right msg

    prop "EDNS" . forAll genEDNSHeader $ \(edns, hdr) -> do
        let eh = EDNSheader edns
            Right m = decode. encode $ DNSMessage hdr eh [] [] [] []
        ednsHeader m `shouldBe` eh

----------------------------------------------------------------

genDNSMessage :: Gen DNSMessage
genDNSMessage =
    DNSMessage <$> genDNSHeader 0x0f <*> makeEDNS <*> listOf genQuestion
               <*> listOf genResourceRecord  <*> listOf genResourceRecord
               <*> listOf genResourceRecord
  where
    makeEDNS :: Gen EDNSheader
    makeEDNS = genBool >>= \t ->
        if t then EDNSheader <$> genEDNS
             else pure NoEDNS


genQuestion :: Gen Question
genQuestion = Question <$> genDomain <*> genTYPE

genTYPE :: Gen TYPE
genTYPE = frequency
    [ (20, elements
            [ A, AAAA, NS, TXT, MX, CNAME, SOA, PTR, SRV, DNAME, OPT, DS, RRSIG
            , NSEC, DNSKEY, NSEC3, NSEC3PARAM, TLSA, CDS, CDNSKEY, CSYNC
            ])
    , (1, toTYPE <$> genWord16)
    ]

genResourceRecord :: Gen ResourceRecord
genResourceRecord = frequency
    [ (8, genRR)
    ]
  where
    genRR = do
      dom <- genDomain
      t <- elements [A, AAAA, NS, TXT, MX, CNAME, SOA, PTR, SRV, DNAME, DS,
                     TLSA, NSEC, NSEC3]
      ResourceRecord dom t classIN <$> genWord32 <*> mkRData dom t

mkRData :: Domain -> TYPE -> Gen RData
mkRData dom typ =
    case typ of
        A -> RD_A <$> genIPv4
        AAAA -> RD_AAAA <$> genIPv6
        NS -> pure $ RD_NS dom
        TXT -> RD_TXT <$> genTextString
        MX -> RD_MX <$> genWord16 <*> genDomain
        CNAME -> pure $ RD_CNAME dom
        SOA -> RD_SOA dom <$> genMailbox <*> genWord32 <*> genWord32 <*> genWord32 <*> genWord32 <*> genWord32
        PTR -> RD_PTR <$> genDomain
        SRV -> RD_SRV <$> genWord16 <*> genWord16 <*> genWord16 <*> genDomain
        DNAME -> RD_DNAME <$> genDomain
        DS -> RD_DS <$> genWord16 <*> genWord8 <*> genWord8 <*> genByteString
        NSEC -> RD_NSEC <$> genDomain <*> genNsecTypes
        NSEC3 -> genNSEC3
        TLSA -> RD_TLSA <$> genWord8 <*> genWord8 <*> genWord8 <*> genByteString
        CAA -> RD_CAA <$> genWord8 <*> (CI.mk <$> genByteString) <*> genByteString

        _ -> pure . RD_TXT $ "Unhandled type " <> BS.pack (show typ)
  where
    genNSEC3 = do
        (alg, hlen)  <- elements [(1,32),(2,64)]
        flgs <- elements [0,1]
        iter <- elements [0..100]
        salt <- elements ["", "AB"]
        hash <- B.pack <$> replicateM hlen genWord8
        RD_NSEC3 alg flgs iter salt hash <$> genNsecTypes
    genTextString = do
        len <- elements [0, 1, 63, 255, 256, 511, 512, 1023, 1024]
        B.pack <$> replicateM len genWord8
    genNsecTypes = do
        ntypes <- elements [0..15]
        types <- sequence $ replicate ntypes $ toTYPE <$> elements [1..1024]
        return $ [ the t |
                   t <- types,
                   then group by (fromTYPE t)
                        using groupWith ]

genIPv4 :: Gen IPv4
genIPv4 = toIPv4 <$> replicateM 4 (fromIntegral <$> genWord8)

genIPv6 :: Gen IPv6
genIPv6 = toIPv6 <$> replicateM 8 (fromIntegral <$> genWord16)

genByteString :: Gen BS.ByteString
genByteString = elements
    [ "", "a", "a.b", "abc", "a.b.c", "a\\.b.c", "\\001.a.b", "\\$.a.b" ]

genMboxString :: Gen BS.ByteString
genMboxString = elements
    [ "", "a", "a@b", "abc", "a@b.c", "first.last@example.org" ]

genDomain :: Gen Domain
genDomain = do
    bs <- genByteString
    pure $ bs <> "."

genMailbox :: Gen Mailbox
genMailbox = do
    bs <- genMboxString
    pure $ bs <> "."

genDNSHeader :: Word16 -> Gen DNSHeader
genDNSHeader maxrc = DNSHeader <$> genWord16 <*> genDNSFlags maxrc

genDNSFlags :: Word16 -> Gen DNSFlags
genDNSFlags maxrc =
  DNSFlags <$> genQorR <*> genOPCODE <*> genBool        <*> genBool
           <*> genBool <*> genBool   <*> genRCODE maxrc <*> genBool <*> genBool

genWord16 :: Gen Word16
genWord16 = arbitrary

genWord32 :: Gen Word32
genWord32 = arbitrary

genWord8 :: Gen Word8
genWord8 = arbitrary

genBool :: Gen Bool
genBool = elements [True, False]

genQorR :: Gen QorR
genQorR = elements [minBound .. maxBound]

genOPCODE :: Gen OPCODE
genOPCODE  = elements [OP_STD, OP_INV, OP_SSR, OP_NOTIFY, OP_UPDATE]

genRCODE :: Word16 -> Gen RCODE
genRCODE maxrc = elements $ map toRCODE [0..maxrc]

genEDNS :: Gen EDNS
genEDNS = do
    vers <- genWord8
    ok <- genBool
    od <- genOData
    us <- elements [minUdpSize..maxUdpSize]
    return $ defaultEDNS {
        ednsVersion  = vers
      , ednsUdpSize  = us
      , ednsDnssecOk = ok
      , ednsOptions  = [od]
      }

genOData :: Gen OData
genOData = oneof
    [ genOD_Unknown
    , genOD_ECS
    ]
  where
    -- | Choose from the range reserved for local use
    -- https://tools.ietf.org/html/rfc6891#section-9
    genOD_Unknown = UnknownOData <$> elements [65001, 65534] <*> genByteString

    -- | Only valid ECS prefixes round-trip, make sure the prefix is
    -- is consistent with the mask.
    genOD_ECS = do
        usev4 <- genBool
        if usev4
        then genFuzzed genIPv4 IPv4 Data.IP.fromIPv4  1 32
        else genFuzzed genIPv6 IPv6 Data.IP.fromIPv6b 2 128
      where
        genFuzzed :: Addr a
                  => Gen a
                  -> (a -> IP)
                  -> (a -> [Int])
                  -> Word16
                  -> Word8
                  -> Gen OData
        genFuzzed gen toIP toBytes fam alen = do
            ip <- gen
            bits1 <- elements [1 .. alen]
            bits2 <- elements [0 .. alen]
            fuzzSrcBits <- genBool
            fuzzScpBits <- genBool
            srcBits <- if not fuzzSrcBits
                       then pure bits1
                       else flip mod alen. (+) bits1 <$> elements [1..alen-1]
            scpBits <- if not fuzzScpBits
                       then pure bits2
                       else elements [alen+1 .. 0xFF]
            let addr  = Data.IP.addr. makeAddrRange ip $ fromIntegral bits1
                bytes = map fromIntegral $ toBytes addr
                len   = (fromIntegral bits1 + 7) `div` 8
                less  = take (len - 1) bytes
                more  = less ++ [0xFF]
            if srcBits == bits1
            then if scpBits == bits2
                 then pure $ OD_ClientSubnet bits1 scpBits $ toIP addr
                 else pure $ OD_ECSgeneric fam bits1 scpBits $ B.pack bytes
            else if srcBits < bits1
                 then pure $ OD_ECSgeneric fam srcBits scpBits $ B.pack more
                 else pure $ OD_ECSgeneric fam srcBits scpBits $ B.pack less

genExtRCODE :: Gen RCODE
genExtRCODE = elements $ map toRCODE [0..4095]

genEDNSHeader :: Gen (EDNS, DNSHeader)
genEDNSHeader = do
    edns <- genEDNS
    hdr <- genDNSHeader 0xF00
    return (edns, hdr)