{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Test.Cloudflare.DNS.Record where

import           Control.Concurrent
import           Control.Monad
import           Crypto.Hash
import qualified Crypto.PubKey.RSA             as RSA
import qualified Crypto.PubKey.RSA.PKCS15      as PKCS15
import           Data.ASN1.OID
import qualified Data.Binary                   as Binary
import qualified Data.Binary.Put               as Binary
import           Data.Bits
import           Data.ByteArray                (convert)
import qualified Data.ByteString.Base64        as B64
import qualified Data.ByteString.Lazy          as BSL
import qualified Data.ByteString.Lazy.Base64   as BSL64
import           Data.Hourglass
import           Data.IP
import qualified Data.List                     as List
import           Data.Text                     (Text)
import qualified Data.Text                     as Text
import qualified Data.Text.Encoding            as TE
import qualified Data.Text.Lazy                as TL
import           Data.Word
import           Data.X509
import           Network.Cloudflare.DNS.Record
import           Network.Cloudflare.Types
import           Network.Cloudflare.Zone
import           System.Hourglass
import           System.Process
import           System.Random
import           Test.Cloudflare.Types
import           Test.Tasty
import           Test.Tasty.HUnit

dnsRecordTests :: CloudflareTestConfig -> TestTree
dnsRecordTests config = testGroup "DNS Record Tests" [
    dnsRecordListTest config
  , dnsRecordDeleteTest config
  , dnsRecordPropagateTest config
  -- , dnsRecordCreateTest config
  ]

-- | Test listing all DNS records for all zones the token has access to
dnsRecordListTest :: CloudflareTestConfig -> TestTree
dnsRecordListTest config =
  testCaseSteps "DNS Record List" $ \step -> do
    step "List all zones"
    zonesResult <- either (failWith "Could not list zones: ") pure =<< listZones authInfo
    forM_ (resultWithInfoResult zonesResult) $ \zone -> do
      step $ "List all records for zone " <> show (zoneName zone)
      listAllRecords authInfo zone
  where
    authInfo = cloudflareTestConfigAuth config

-- | Test deleting all DNS records for all zones the token has access to
dnsRecordDeleteTest :: CloudflareTestConfig -> TestTree
dnsRecordDeleteTest config =
  testCaseSteps "DNS Record Delete" $ \step -> do
    step "List all zones"
    zonesResult <- either (failWith "Could not list zones: ") pure =<< listZones authInfo
    forM_ (resultWithInfoResult zonesResult) $ \zone -> do
      deleteAllRecords step authInfo zone
  where
    authInfo = cloudflareTestConfigAuth config

-- | Test creating a DNS record for all zones the token has access to
-- deletes all records for the zone before creating the record
-- and deletes the record after checking that it's correct
-- checks the result of the create operation
dnsRecordCreateTest :: CloudflareTestConfig -> TestTree
dnsRecordCreateTest config =
  testCaseSteps "DNS Record Create" $ \step -> do
    step "Create test certificate"
    -- ((pubKey, _privKey), signed) <- genTestCertificate
    step "List all zones"
    zonesResult <- either (failWith "Could not list zones: ") pure =<< listZones authInfo
    forM_ (resultWithInfoResult zonesResult) $ \zone -> do
      deleteAllRecords step authInfo zone
      -- A/AAAA records cannot exist with CNAME records
      -- NS records cannot exist with any other record
      createRecord step config zone (CreateDNSRecord (ARecord "192.0.2.0") (cloudflareTestConfigDomain config) False (Just "A test comment ARecord") ["A", "tag1"] 60)
      createRecord step config zone (CreateDNSRecord (AAAARecord "2001:db8:ffff:ffff:ffff:ffff:ffff:ffff") (cloudflareTestConfigDomain config) False (Just "A test comment AAAARecord") ["AAAA", "tag2"] 60)
      createRecord step config zone (CreateDNSRecord (CAARecord 0 "issue" (cloudflareTestConfigDomain config)) (cloudflareTestConfigDomain config) False (Just "A test comment AAAARecord") ["AAAA", "tag3"] 60)
      -- createRecord step config zone (CreateDNSRecord (CERTRecord 5 (x509EncodeCertificate signed) (keyTagRSA pubKey) 1) (cloudflareTestConfigDomain config) False (Just "A test comment CERTRecord") ["CERT", "tag4"] 60)
      -- let dnskey = DNSKEYRecord 5 0 3 (encodeRSAPublicKey pubKey)
      -- createRecord step config zone (CreateDNSRecord dnskey (cloudflareTestConfigDomain config) False (Just "A test comment DNSKEY") ["DNSKEY", "tag5"] 60)
      -- let ds = mkDSRecord (cloudflareTestConfigDomain config) dnskey
      -- createRecord step config zone (CreateDNSRecord ds (cloudflareTestConfigDomain config) False (Just "A test comment DS") ["DS", "tag6"] 60)
      -- createRecord step config zone (CreateDNSRecord (HTTPSRecord 0 (cloudflareTestConfigDomain config) "alpn=\"h3,h2\" ipv4hint=\"127.0.0.1\" ipv6hint=\"::1\"") (cloudflareTestConfigDomain config) False (Just "A test comment HTTPS") ["HTTPS", "tag7"] 60)
      createRecord step config zone (CreateDNSRecord (LOCRecord 35 40 "N" 46 57.284 73 "W" 57 56.177 5 5 5) (cloudflareTestConfigDomain config) False (Just "A test comment LOCRecord") ["LOC", "tag8"] 60)
      -- createRecord step config zone (CreateDNSRecord (NAPTRRecord "u" 100 10 "!^.*$!mailto:information@foo.se!i" "." "smtp+E2U" ) (cloudflareTestConfigDomain config) False (Just "A test comment LOCRecord") ["LOC", "tag9"] 60)
      -- createRecord step config zone (CreateDNSRecord (PTRRecord "192.0.2.0") (cloudflareTestConfigDomain config) False (Just "A test comment LOCRecord") ["LOC", "tag10"] 60)
      -- SMIMEARecord
      -- SRVRecord
      createRecord step config zone (CreateDNSRecord (SRVRecord "example.com" 5060 10 "_TCP" "_SIP" "server.example.com" 20) (cloudflareTestConfigDomain config) False (Just "A test comment LOCRecord") ["LOC", "tag13"] 60)
      -- SSHFPRecord
      -- SVCBRecord
      -- TLSARecord
      createRecord step config zone (CreateDNSRecord (TXTRecord "Hello World") (cloudflareTestConfigDomain config) False (Just "A test comment TXTRecord") ["TXT", "tag17"] 60)
      -- URIRecord
      createRecord step config zone (CreateDNSRecord (URIRecord "http://example.com/example.html" 0) (cloudflareTestConfigDomain config) False (Just "A test comment URI Record") ["URI", "tag18"] 60)

      -- MXRecord
      listAllRecords authInfo zone
      -- CNAMERecord
      -- NSRecord

  where
    authInfo = cloudflareTestConfigAuth config

-- | rfc4034 Section 5.1.4
mkDSRecord :: Text -> DNSKEYRecord -> DSRecord
mkDSRecord owner (DNSKEYRecord algorithm flags protocol publicKey) =
  DSRecord 1 (TE.decodeUtf8 $ convert digest) digestType keyTag
  where
    digestType = 1 -- SHA-1
    digest = hashWith SHA1 $ BSL.toStrict $ mconcat [
        BSL.fromStrict $ TE.encodeUtf8 owner
      , Binary.runPut $ Binary.putWord8 (fromIntegral flags)
      , Binary.runPut $ Binary.putWord8 (fromIntegral protocol)
      , Binary.runPut $ Binary.putWord16be (fromIntegral algorithm)
      , Binary.runPut $ Binary.putByteString $ TE.encodeUtf8 publicKey
      ]
    keyTag = 1 -- TODO

-- | RFC3110 Section 2
encodeRSAPublicKey :: RSA.PublicKey -> Text
encodeRSAPublicKey pubKey = TL.toStrict $ BSL64.encodeBase64 $ Binary.runPut $ do
  Binary.putWord8 (0 :: Word8)
  Binary.putWord16be (fromIntegral $ RSA.public_size pubKey)
  Binary.put $ toBytes $ RSA.public_e pubKey
  Binary.put $ toBytes $ RSA.public_n pubKey
  where
    toBytes :: Integer -> [Word8]
    toBytes = dropWhile (== 0) . reverse . go
      where
        go 0 = []
        go i = fromIntegral (i .&. 0xff) : go (shiftR i 8)
    
-- | Section 3.1 of RFC4034
-- The certificate MUST be encoded as described in Section

x509EncodeCertificate :: SignedExact Certificate -> Text
x509EncodeCertificate = B64.encodeBase64 . encodeSignedObject

genTestCertificate :: IO ((RSA.PublicKey, RSA.PrivateKey), SignedExact Certificate)
genTestCertificate = do
  now <- dateCurrent
  let validity = (now, timeAdd now (Hours 24))
  keyPair@(pubKey, privKey) <- RSA.generate 512 65537
  let certificate = Certificate version serial signatureAlg issuerDN validity subjectDN (PubKeyRSA pubKey) extensions
      (signed, _) = objectToSignedExact (sign privKey) certificate 
  pure (keyPair, signed)
  where
    version = 0
    serial = 1
    signatureAlg = SignatureALG HashSHA1 PubKeyALG_RSA
    issuerDN = DistinguishedName []
    subjectDN = DistinguishedName [
        (getObjectID DnCommonName, "test")
      , (getObjectID DnCountry, "US")
      , (getObjectID DnOrganization, "test")
      ]
    extensions = Extensions Nothing
    sign priv b = (sig, SignatureALG HashSHA1 (privkeyToAlg (PrivKeyRSA priv)), ())
      where sig = either (error . ("Signature error: " <> ) . show) id $ PKCS15.sign Nothing (Just SHA1) priv b

-- | Section B.1 of RFC4034
-- the key tag is defined to be the most significant 16 bits of
-- the least significant 24 bits in the public key modulus
keyTagRSA :: RSA.PublicKey -> Word16
keyTagRSA pubKey = fromIntegral $ shiftR (RSA.public_n pubKey) 8

-- | Delete all DNS records for a zone. Checks the results of the delete operation
deleteAllRecords :: (String -> IO ())
                 -> CloudflareAuth
                 -> Zone
                 -> Assertion
deleteAllRecords step authInfo zone = do
  beforeDeleteRecordsResult <- either (failWith "Could not list records: ") pure =<< listDNSRecords authInfo (zoneId zone)
  forM_ (resultWithInfoResult beforeDeleteRecordsResult) $ \record -> do
    step $ "Delete record: " <> show (withDNSEntry dnsRecordEntryType record)
    deleteResult <- either (failWith "Could not delete record: ") pure =<< deleteDNSRecord authInfo (zoneId zone) (withDNSEntry dnsRecordEntryId record)
    resultOnlyResult deleteResult @?= (DNSIdentifier (withDNSEntry dnsRecordEntryId record))
  afterDeleteRecordsResult <- either (failWith "Could not list records: ") pure =<< listDNSRecords authInfo (zoneId zone)
  step "Check Delete result"
  checkResultWithInfo afterDeleteRecordsResult
  resultWithInfoResult afterDeleteRecordsResult @?= []

listAllRecords :: CloudflareAuth
               -> Zone
               -> Assertion
listAllRecords authInfo zone = do
    recordsResult <- either (failWith "Could not list records: ") pure =<< listDNSRecords authInfo (zoneId zone)
    checkResultWithInfo recordsResult

createRecord :: (Eq record, Show record, DNSRecord record)
             => (String -> IO ())
             -> CloudflareTestConfig
             -> Zone
             -> CreateDNSRecord record
             -> Assertion
createRecord step config zone record' = do
  step $ "Create record: " <> show record
  createResult <- either (failWith "Could not create record: ") pure =<< createDNSRecord authInfo (zoneId zone) record
  checkResultResponse createResult
  checkCreatedRecord record (resultResult createResult)
  where
    record = createDNSRecordFeatures config record'
    authInfo = cloudflareTestConfigAuth config

-- | Filter out pro features if the config has them disabled
createDNSRecordFeatures :: CloudflareTestConfig -> CreateDNSRecord record -> CreateDNSRecord record
createDNSRecordFeatures config record =
  if cloudflareTestConfigProFeatures config
    then record
    else record {
            createDNSRecordTags = []
          }

checkCreatedRecord :: (Eq record, Show record)
                   => CreateDNSRecord record
                   -> DNSRecordEntry record
                   -> Assertion
checkCreatedRecord record e = do
  assertEqual "Created record data" (createDNSRecordData record) (dnsRecordEntryData e)
  assertEqual "Created record tags" (createDNSRecordTags record) (dnsRecordEntryTags e)
  assertEqual "Created record comment" (createDNSRecordComment record) (dnsRecordEntryComment e)
  assertEqual "Created record TTL" (createDNSRecordTTL record) (dnsRecordEntryTTL e)

checkResultWithInfo :: ResultWithInfo a -> Assertion
checkResultWithInfo info = do
  assertEqual "Result errors" [] (resultWithInfoErrors info)
  assertEqual "Result messages" [] (resultWithInfoMessages info)
  assertBool "Result success" (resultWithInfoSuccess info)

checkResultResponse :: ResultResponse a -> Assertion
checkResultResponse response = do
  assertEqual "Result errors" [] (resultErrors response)
  assertEqual "Result messages" [] (resultMessages response)
  assertBool "Result success" (resultSuccess response)

findZoneForDomain :: CloudflareAuth -> Text -> IO Zone
findZoneForDomain authInfo domain = do
  zonesResult <- either (failWith "Could not list zones: ") pure =<< listZones authInfo
  let findZone = List.find (\z -> zoneName z == domain)
  maybe (assertFailure $ "Could not find zone for " <> Text.unpack domain) pure $ findZone (resultWithInfoResult zonesResult)
  

-- Tests that cloudflare api changes propagate to the dns servers
dnsRecordPropagateTest :: CloudflareTestConfig -> TestTree
dnsRecordPropagateTest config =
  testCaseSteps "DNS Record Create and Propagate" $ \step -> do
    step "Create test certificate"
    -- ((pubKey, _privKey), signed) <- genTestCertificate
    step "List all zones"
    zone <- findZoneForDomain authInfo domain
    createAndWaitForPropagate config step zone aRecordTest
    createAndWaitForPropagate config step zone aaaaRecordTest
    createAndWaitForPropagate config step zone mxRecordTest
    createAndWaitForPropagate config step zone locRecordTest
    createAndWaitForPropagate config step zone txtRecordTest
    createAndWaitForPropagate config step zone caaRecordTest
    createAndWaitForPropagate config step zone uriRecordTest
  where
    domain = cloudflareTestConfigDomain config
    authInfo = cloudflareTestConfigAuth config
    aRecordTest :: DNSPropagateTestVector ARecord
    aRecordTest = DNSPropagateTestVector {
        dnsPropagateTestVectorCreate = do
          end <- randomRIO (0, 255)
          let ip = toIPv4 [192, 0, 2, end]
          pure $ CreateDNSRecord (ARecord $ Text.pack $ show ip) (cloudflareTestConfigDomain config) False (Just "A test comment ARecord") ["A", "tag1"] 60
      , dnsPropagateTestVectorExpected = \r -> [Text.unpack $ aRecordContent $ createDNSRecordData r]
      , dnsPropagateTestVectorDelvArgs = ["-i"]
    }
    aaaaRecordTest :: DNSPropagateTestVector AAAARecord
    aaaaRecordTest = DNSPropagateTestVector {
        dnsPropagateTestVectorCreate = do
          end <- randomRIO (0, 65535)
          let ip = toIPv6 [0x2001, 0xdb8, 0xffff, 0xffff, 0xffff, 0xffff, 0xffff, end]
          pure $ CreateDNSRecord (AAAARecord $ Text.pack $ show ip) (cloudflareTestConfigDomain config) False (Just "A test comment AAAARecord") ["AAAA", "tag2"] 60
      , dnsPropagateTestVectorExpected = \r -> [Text.unpack $ aaaaRecordContent $ createDNSRecordData r]
      , dnsPropagateTestVectorDelvArgs = ["-i"]
    }
    mxRecordTest :: DNSPropagateTestVector MXRecord
    mxRecordTest = DNSPropagateTestVector {
        dnsPropagateTestVectorCreate = do
          end :: Int <- randomRIO (0, 65535)
          let mxName = "test" <> show end <> "." <> Text.unpack (cloudflareTestConfigDomain config)
          pure $ CreateDNSRecord (MXRecord (Text.pack mxName) 0) (cloudflareTestConfigDomain config) False (Just "A test comment MXRecord") ["MX", "tag3"] 60
      , dnsPropagateTestVectorExpected = \r ->
          let record = createDNSRecordData r
          in [show (mxRecordPriority record) <> " " <> Text.unpack (mxRecordContent record) <> "."]
      , dnsPropagateTestVectorDelvArgs = ["-i"]
    }
    locRecordTest :: DNSPropagateTestVector LOCRecord
    locRecordTest = DNSPropagateTestVector {
        dnsPropagateTestVectorCreate = do
          lat :: Int <- randomRIO (0, 59)
          long :: Int <- randomRIO (0, 59)
          pure $ CreateDNSRecord (LOCRecord 35 lat "N" 46 57.284 73 "W" long 56.177 5 5 5) (cloudflareTestConfigDomain config) False (Just "A test comment LOCRecord") ["LOC", "tag4"] 60
      , dnsPropagateTestVectorExpected = \r ->
        let record = createDNSRecordData r
        in [List.intercalate " " [
              show (locRecordLatDegrees record), show (locRecordLatMinutes record), show (locRecordLatSeconds record), Text.unpack (locRecordLatDirection record)
              , show (locRecordLongDegrees record), show (locRecordLongMinutes record), show (locRecordLongSeconds record), Text.unpack (locRecordLongDirection record)
              , show (locRecordAltitude record) <> ".00m", show (locRecordSize record) <> "m", show (locRecordPrecisionHoriz record) <> "m", show (locRecordPrecisionVert record) <> "m"]]
      , dnsPropagateTestVectorDelvArgs = ["-i"]
    }
    txtRecordTest :: DNSPropagateTestVector TXTRecord
    txtRecordTest = DNSPropagateTestVector {
        dnsPropagateTestVectorCreate = do
          let
            txtWords :: [String]
            txtWords = ["Hello", "world!", "foo", "bar", "baz", "qux"]
          len <- randomRIO (1, length txtWords)
          let content = List.intercalate " " $ take len txtWords
          
          pure $ CreateDNSRecord (TXTRecord $ Text.pack content) (cloudflareTestConfigDomain config) False (Just "A test comment TXTRecord") ["TXT", "tag5"] 60
      , dnsPropagateTestVectorExpected = \r -> [show $ txtRecordContent $ createDNSRecordData r]
      , dnsPropagateTestVectorDelvArgs = ["-i"]
    }
    caaRecordTest :: DNSPropagateTestVector CAARecord
    caaRecordTest = DNSPropagateTestVector {
        dnsPropagateTestVectorCreate = do
          nonce :: Int <- randomRIO (1, 1000)
          pure $ CreateDNSRecord (CAARecord 0 "issue" ("ca" <> Text.pack (show nonce) <> ".example.net")) (cloudflareTestConfigDomain config) False (Just "A test comment CAARecord") ["CAA", "tag5"] 60
      , dnsPropagateTestVectorExpected = \r -> [List.intercalate " " [
            show $ caaRecordFlags $ createDNSRecordData r
          , Text.unpack $ caaRecordTag $ createDNSRecordData r
          , show $ caaRecordValue $ createDNSRecordData r
          ]]
      , dnsPropagateTestVectorDelvArgs = ["-i"]
    }
    uriRecordTest :: DNSPropagateTestVector URIRecord
    uriRecordTest = DNSPropagateTestVector {
        dnsPropagateTestVectorCreate = do
          nonce :: Int <- randomRIO (1, 1000)
          pure $ CreateDNSRecord (URIRecord ("http://example.com/" <> Text.pack (show nonce)) 10) (cloudflareTestConfigDomain config) False (Just "A test comment CAARecord") ["CAA", "tag5"] 60
      , dnsPropagateTestVectorExpected = \r -> [List.intercalate " " [
            show $ uriRecordPriority $ createDNSRecordData r
          , Text.unpack $ uriRecordContent $ createDNSRecordData r
          ]]
      , dnsPropagateTestVectorDelvArgs = ["-i"]
    }
    -- untested records: cert, https, naptr, ptr, smimea, srv, sshfp, svcb tlsa, cname
    -- untested dnssec records: dnskey, ds
    -- requires a third domain: ns

    -- This test has to be ran separately than the A/AAAA tests
    -- cnameRecordTest :: DNSPropagateTestVector CNAMERecord
    -- cnameRecordTest = DNSPropagateTestVector {
    --     dnsPropagateTestVectorCreate = do
    --       end :: Int <- randomRIO (0, 65535)
    --       let mxName = "test" <> show end <> "." <> Text.unpack (cloudflareTestConfigDomain config)
    --       pure $ CreateDNSRecord (CNAMERecord (Text.pack mxName)) (cloudflareTestConfigDomain config) False (Just "A test comment CNAMERecord") ["MX", "tag6"] 60
    --   , dnsPropagateTestVectorExpected = \r ->
    --       let record = createDNSRecordData r
    --       in [Text.unpack (cnameRecordContent record) <> "."]
    --   , dnsPropagateTestVectorDelvArgs = ["-i"]
    -- }

createAndWaitForPropagate :: (Eq record, Show record, DNSRecord record)
                          => CloudflareTestConfig
                          -> (String -> IO ())
                          -> Zone
                          -> DNSPropagateTestVector record
                          -> Assertion
createAndWaitForPropagate config step zone test = do
  create <- dnsPropagateTestVectorCreate test
  createRecord step config zone create
  waitForDelvPropagation "1.1.1.1" delvTarget (recordType create) delvArgs 600 (dnsPropagateTestVectorExpected test create)
  where
    delvTarget = Text.unpack $ cloudflareTestConfigDomain config
    delvArgs = dnsPropagateTestVectorDelvArgs test
    recordType = Text.unpack . dnsRecordType . createDNSRecordData

requiresDNSSec :: DNSRecord record => record -> Bool
requiresDNSSec record =
  case dnsRecordType record of
    "DLV" -> True
    "DNSKEY" -> True
    "DS" -> True
    "NSEC" -> True
    "NSEC3" -> True
    "RSIG" -> True
    _ -> False

data DNSPropagateTestVector record = DNSPropagateTestVector {
  dnsPropagateTestVectorCreate :: IO (CreateDNSRecord record)
, dnsPropagateTestVectorExpected :: CreateDNSRecord record -> [String]
, dnsPropagateTestVectorDelvArgs :: [String]
}

waitForDelvPropagation :: String -- ^ Resolver IP address
                       -> String -- ^ Domain name
                       -> String -- ^ Record type
                       -> [String] -- ^ Extra arguments to delv
                       -> Int -- ^ Timeout in seconds
                       -> [String] -- ^ Expected short output from delv
                       -> Assertion
waitForDelvPropagation resolver domain record args seconds expected = do
  go 0 >>= \case
    Left res -> assertFailure $ "Timed out waiting for propagation of record " <> record <> " for domain " <> domain <> " to " <> show expected <> " but got " <> show res
    Right () -> pure ()
  where
    go n = do
      threadDelay 1000000
      actual <- delvFor resolver domain record args
      if actual == expected
        then pure $ Right ()
        else if (n < seconds)
          then go (n + 1)
          else pure $ Left actual

delvFor :: String -- ^ Resolver IP address
        -> String -- ^ Domain name
        -> String -- ^ Record type
        -> [String] -- ^ Extra arguments to delv
        -> IO [String]
delvFor resolver domain record args = do
  let command = shell $ "delv " <> List.intercalate " " (["@" <> resolver, record, domain, "+short"] <> args)
  p <- readCreateProcess (command { std_err = CreatePipe } ) ""
  pure $ lines p