From 0807113e6c445171b62a0a18a53a0bd09f950af5 Mon Sep 17 00:00:00 2001
From: mzack <marco.rivoli.nvh@gmail.com>
Date: Fri, 12 Apr 2024 00:02:43 +0200
Subject: [PATCH] adding more query types test

---
 cmd/integration-test/dns.go                    | 18 +++++++++++++++---
 .../protocols/dns/{basic.yaml => a.yaml}       |  6 +++---
 integration_tests/protocols/dns/aaaa.yaml      | 17 +++++++++++++++++
 3 files changed, 35 insertions(+), 6 deletions(-)
 rename integration_tests/protocols/dns/{basic.yaml => a.yaml} (69%)
 create mode 100644 integration_tests/protocols/dns/aaaa.yaml

diff --git a/cmd/integration-test/dns.go b/cmd/integration-test/dns.go
index 982aab98b..9f1883716 100644
--- a/cmd/integration-test/dns.go
+++ b/cmd/integration-test/dns.go
@@ -5,7 +5,8 @@ import (
 )
 
 var dnsTestCases = []TestCaseInfo{
-	{Path: "protocols/dns/basic.yaml", TestCase: &dnsBasic{}},
+	{Path: "protocols/dns/a.yaml", TestCase: &dnsA{}},
+	{Path: "protocols/dns/aaaa.yaml", TestCase: &dnsAAAA{}},
 	{Path: "protocols/dns/ptr.yaml", TestCase: &dnsPtr{}},
 	{Path: "protocols/dns/caa.yaml", TestCase: &dnsCAA{}},
 	{Path: "protocols/dns/tlsa.yaml", TestCase: &dnsTLSA{}},
@@ -14,10 +15,21 @@ var dnsTestCases = []TestCaseInfo{
 	{Path: "protocols/dns/dsl-matcher-variable.yaml", TestCase: &dnsDSLMatcherVariable{}},
 }
 
-type dnsBasic struct{}
+type dnsA struct{}
 
 // Execute executes a test case and returns an error if occurred
-func (h *dnsBasic) Execute(filePath string) error {
+func (h *dnsA) Execute(filePath string) error {
+	results, err := testutils.RunNucleiTemplateAndGetResults(filePath, "one.one.one.one", debug)
+	if err != nil {
+		return err
+	}
+	return expectResultsCount(results, 1)
+}
+
+type dnsAAAA struct{}
+
+// Execute executes a test case and returns an error if occurred
+func (h *dnsAAAA) Execute(filePath string) error {
 	results, err := testutils.RunNucleiTemplateAndGetResults(filePath, "one.one.one.one", debug)
 	if err != nil {
 		return err
diff --git a/integration_tests/protocols/dns/basic.yaml b/integration_tests/protocols/dns/a.yaml
similarity index 69%
rename from integration_tests/protocols/dns/basic.yaml
rename to integration_tests/protocols/dns/a.yaml
index b0dbdf3b0..538974626 100644
--- a/integration_tests/protocols/dns/basic.yaml
+++ b/integration_tests/protocols/dns/a.yaml
@@ -1,7 +1,7 @@
-id: basic-dns-example
+id: basic-dns-a-example
 
 info:
-  name: Test DNS Template
+  name: Test DNS A Query Template
   author: pdteam
   severity: info
 
@@ -14,4 +14,4 @@ dns:
     matchers:
       - type: word
         words:
-          - "1.1.1.1"
\ No newline at end of file
+          - "1.1.1.1"
diff --git a/integration_tests/protocols/dns/aaaa.yaml b/integration_tests/protocols/dns/aaaa.yaml
new file mode 100644
index 000000000..3df3293a9
--- /dev/null
+++ b/integration_tests/protocols/dns/aaaa.yaml
@@ -0,0 +1,17 @@
+id: basic-dns-aaaa-example
+
+info:
+  name: Test DNS AAAA Query Template
+  author: pdteam
+  severity: info
+
+dns:
+  - name: "{{FQDN}}"
+    type: AAAA
+    class: inet
+    recursion: true
+    retries: 3
+    matchers:
+      - type: word
+        words:
+          - "2606:4700:4700::1001"