From ec230546d8aeb7a732b609eb4b0293d08d5aa077 Mon Sep 17 00:00:00 2001
From: Alden Hilton <106177711+adhilto@users.noreply.github.com>
Date: Mon, 29 Jan 2024 12:33:22 -0800
Subject: [PATCH] Correct DNS bug (#160)

---
 Testing/RegoTests/gmail/gmail02_test.rego |  12 +-
 Testing/RegoTests/gmail/gmail03_test.rego |  44 +--
 Testing/RegoTests/gmail/gmail04_test.rego | 320 +++-------------------
 rego/Gmail.rego                           |   4 +-
 4 files changed, 49 insertions(+), 331 deletions(-)

diff --git a/Testing/RegoTests/gmail/gmail02_test.rego b/Testing/RegoTests/gmail/gmail02_test.rego
index 3b7eb1f3..1b32427b 100644
--- a/Testing/RegoTests/gmail/gmail02_test.rego
+++ b/Testing/RegoTests/gmail/gmail02_test.rego
@@ -14,7 +14,8 @@ test_DKIM_Correct_V1 if {
                 "domain": "test.name",
                 "rdata": ["v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt"]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -37,7 +38,8 @@ test_DKIM_Correct_V2 if {
                 "domain": "test2.name",
                 "rdata": ["v=DKIM1;"]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -60,7 +62,8 @@ test_DKIM_Incorrect_V1 if {
                 "domain": "test2.name",
                 "rdata": []
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -79,7 +82,8 @@ test_DKIM_Incorrect_V2 if {
                 "domain": "test.name",
                 "rdata": []
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
diff --git a/Testing/RegoTests/gmail/gmail03_test.rego b/Testing/RegoTests/gmail/gmail03_test.rego
index 199c8ac6..2df5f349 100644
--- a/Testing/RegoTests/gmail/gmail03_test.rego
+++ b/Testing/RegoTests/gmail/gmail03_test.rego
@@ -9,18 +9,13 @@ test_SPF_Correct_V1 if {
     # Test SPF when there's only one domain
     PolicyId := "GWS.GMAIL.3.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": ["v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+"]
-            }
-        ],
         "spf_records": [
             {
                 "domain": "test.name",
                 "rdata": ["v=spf1 include:_spf.google.com ~all"]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -34,16 +29,6 @@ test_SPF_Correct_V2 if {
     # Test SPF when there's multiple domains
     PolicyId := "GWS.GMAIL.3.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": ["v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+"]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "spf_records": [
             {
                 "domain": "test1.name",
@@ -53,7 +38,8 @@ test_SPF_Correct_V2 if {
                 "domain": "test2.name",
                 "rdata": ["v=spf1 "]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -67,16 +53,6 @@ test_SPF_Incorrect_V1 if {
     # Test SPF when there's multiple domains and only one is correct
     PolicyId := "GWS.GMAIL.3.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": ["v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+"]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "spf_records": [
             {
                 "domain": "test1.name",
@@ -86,7 +62,8 @@ test_SPF_Incorrect_V1 if {
                 "domain": "test2.name",
                 "rdata": []
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -100,18 +77,13 @@ test_SPF_Incorrect_V2 if {
     # Test SPF when there's only one domain and it's wrong
     PolicyId := "GWS.GMAIL.3.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": ["v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+"]
-            }
-        ],
         "spf_records": [
             {
                 "domain": "test.name",
                 "rdata": []
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
diff --git a/Testing/RegoTests/gmail/gmail04_test.rego b/Testing/RegoTests/gmail/gmail04_test.rego
index 392050c0..8a0a6714 100644
--- a/Testing/RegoTests/gmail/gmail04_test.rego
+++ b/Testing/RegoTests/gmail/gmail04_test.rego
@@ -9,21 +9,6 @@ test_DMARC_Correct_V1 if {
     # Test DMARC when there's only one domain
     PolicyId := "GWS.GMAIL.4.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
@@ -31,7 +16,8 @@ test_DMARC_Correct_V1 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -45,25 +31,6 @@ test_DMARC_Correct_V2 if {
     # Test DMARC when there's multiple domains
     PolicyId := "GWS.GMAIL.4.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -77,7 +44,8 @@ test_DMARC_Correct_V2 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -91,25 +59,6 @@ test_DMARC_Incorrect_V1 if {
     # Test DMARC when there's multiple domains and only one is correct
     PolicyId := "GWS.GMAIL.4.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -121,7 +70,8 @@ test_DMARC_Incorrect_V1 if {
                 "domain": "test2.name",
                 "rdata": []
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -135,27 +85,13 @@ test_DMARC_Incorrect_V2 if {
     # Test DMARC when there's only one domain and it's wrong
     PolicyId := "GWS.GMAIL.4.1v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
                 "rdata": []
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -173,21 +109,6 @@ test_DMARCMessageReject_Correct_V1 if {
     # Test DMARC when there's only one domain
     PolicyId := "GWS.GMAIL.4.2v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
@@ -195,7 +116,8 @@ test_DMARCMessageReject_Correct_V1 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -209,25 +131,6 @@ test_DMARCMessageReject_Correct_V2 if {
     # Test DMARC when there's multiple domains
     PolicyId := "GWS.GMAIL.4.2v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -241,7 +144,8 @@ test_DMARCMessageReject_Correct_V2 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -255,25 +159,6 @@ test_DMARCMessageReject_Incorrect_V1 if {
     # Test DMARC when there's multiple domains and only one is correct
     PolicyId := "GWS.GMAIL.4.2v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -285,7 +170,8 @@ test_DMARCMessageReject_Incorrect_V1 if {
                 "domain": "test2.name",
                 "rdata": ["v=DMARC1; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -299,27 +185,13 @@ test_DMARCMessageReject_Incorrect_V2 if {
     # Test DMARC when there's only one domain and it's wrong
     PolicyId := "GWS.GMAIL.4.2v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
                 "rdata": ["v=DMARC1; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -337,21 +209,6 @@ test_DMARCAggregateReports_Correct_V1 if {
     # Test DMARC when there's only one domain
     PolicyId := "GWS.GMAIL.4.3v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
@@ -359,7 +216,8 @@ test_DMARCAggregateReports_Correct_V1 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -373,25 +231,6 @@ test_DMARCAggregateReports_Correct_V2 if {
     # Test DMARC when there's multiple domains
     PolicyId := "GWS.GMAIL.4.3v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -405,7 +244,8 @@ test_DMARCAggregateReports_Correct_V2 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -419,25 +259,6 @@ test_DMARCAggregateReports_Incorrect_V1 if {
     # Test DMARC when there's multiple domains and only one is correct
     PolicyId := "GWS.GMAIL.4.3v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -449,7 +270,8 @@ test_DMARCAggregateReports_Incorrect_V1 if {
                 "domain": "test2.name",
                 "rdata": ["v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov"]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -463,27 +285,13 @@ test_DMARCAggregateReports_Incorrect_V2 if {
     # Test DMARC when there's only one domain and it's wrong
     PolicyId := "GWS.GMAIL.4.3v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
                 "rdata": ["v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov"]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -501,21 +309,6 @@ test_DMARCAgencyPOC_Correct_V1 if {
     # Test DMARC when there's only one domain
     PolicyId := "GWS.GMAIL.4.4v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
@@ -523,7 +316,8 @@ test_DMARCAgencyPOC_Correct_V1 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -537,25 +331,6 @@ test_DMARCAgencyPOC_Correct_V2 if {
     # Test DMARC when there's multiple domains
     PolicyId := "GWS.GMAIL.4.4v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -569,7 +344,8 @@ test_DMARCAgencyPOC_Correct_V2 if {
                     "v=DMARC1; p=reject; pct=100; rua=mailto:DMARC@hq.dhs.gov, mailto:reports@dmarc.cyber.dhs.gov"
                 ]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -583,25 +359,6 @@ test_DMARCAgencyPOC_Incorrect_V1 if {
     # Test DMARC when there's multiple domains and only one is correct
     PolicyId := "GWS.GMAIL.4.4v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test1.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            },
-            {
-                "domain": "test2.name",
-                "rdata": ["v=DKIM1;"]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test1.name",
@@ -613,7 +370,8 @@ test_DMARCAgencyPOC_Incorrect_V1 if {
                 "domain": "test2.name",
                 "rdata": ["v=DMARC1; p=reject; pct=100; mailto:reports@dmarc.cyber.dhs.gov"]
             }
-        ]
+        ],
+        "domains": ["test1.name", "test2.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
@@ -627,27 +385,13 @@ test_DMARCAgencyPOC_Incorrect_V2 if {
     # Test DMARC when there's only one domain and it's wrong
     PolicyId := "GWS.GMAIL.4.4v0.1"
     Output := tests with input as {
-        "dkim_records": [
-            {
-                "domain": "test.name",
-                "rdata": [
-                    concat("", [
-                    "v=DKIM1; k=rsa; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAlaknWsKvtbTLAxtWSF5sDt+",
-                    "zvQhTXhT7V2QTnhPGlVXotXxL4VscG5cSnWus8rS4itN9ItxtzompUVRZ14b6hO1C+",
-                    "pxYAcl8Zaj6wsjE2vmEAmLHeXjj9EHMzrhfay2A02MJHReszokyLKBm+",
-                    "OZ7F4SNWP4SCazXkouOeATNrcIPUZxBV769ewx6ClumvOeHA\" ",
-                    "\"qC77VxJieBg+7LaORrm23DMtWqdkMUWB/wmfCHO333/u6bY21eCMgiP/f",
-                    "+jSiylKDdY5kERpRU0NiIxlTGUhqROJESnxNUTqbK69CTAOYR6qhwJeT4OCsuE1zu6gxANmZMClIMiM2SuntXwNswb4QIDAQAB"
-                    ])
-                ]
-            }
-        ],
         "dmarc_records": [
             {
                 "domain": "test.name",
                 "rdata": ["v=DMARC1; p=reject; pct=100; mailto:reports@dmarc.cyber.dhs.gov"]
             }
-        ]
+        ],
+        "domains": ["test.name"]
     }
 
     RuleOutput := [Result | some Result in Output; Result.PolicyId == PolicyId]
diff --git a/rego/Gmail.rego b/rego/Gmail.rego
index 4574210e..7f189a6c 100644
--- a/rego/Gmail.rego
+++ b/rego/Gmail.rego
@@ -19,9 +19,7 @@ ReportDetailsArray(Status, Array1, Array2) := Detail if {
     Detail := Description(Fraction, " agency domain(s) found in violation: ", String)
 }
 
-AllDomains contains Domain.domain if {
-    some Domain in input.dkim_records
-}
+AllDomains := {Domain | some Domain in input.domains}
 
 LogEvents := utils.GetEvents("gmail_logs")