From 0b45ac790d8e16a17549ce8be946a9efe00207ce Mon Sep 17 00:00:00 2001
From: Bozhan Liang <lbz@rivensbane.com>
Date: Fri, 13 Dec 2024 17:13:46 +0800
Subject: [PATCH] Add GeoIPQueryAttribute to LUA

---
 pdns/lua-record.cc                          |  1 +
 regression-tests.auth-py/test_LuaRecords.py | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/pdns/lua-record.cc b/pdns/lua-record.cc
index 1f2344a881b5..a3f64a761b64 100644
--- a/pdns/lua-record.cc
+++ b/pdns/lua-record.cc
@@ -1245,6 +1245,7 @@ static void setupLuaRecords(LuaContext& lua) // NOLINT(readability-function-cogn
   lua.writeFunction("geoiplookup", [](const string &ip, const GeoIPInterface::GeoIPQueryAttribute attr) {
     return getGeo(ip, attr);
   });
+  lua.writeVariable("GeoIPQueryAttribute", std::unordered_map<std::string, int>{{"ASn", GeoIPInterface::GeoIPQueryAttribute::ASn}, {"City", GeoIPInterface::GeoIPQueryAttribute::City}, {"Continent", GeoIPInterface::GeoIPQueryAttribute::Continent}, {"Country", GeoIPInterface::GeoIPQueryAttribute::Country}, {"Country2", GeoIPInterface::GeoIPQueryAttribute::Country2}, {"Name", GeoIPInterface::GeoIPQueryAttribute::Name}, {"Region", GeoIPInterface::GeoIPQueryAttribute::Region}, {"Location", GeoIPInterface::GeoIPQueryAttribute::Location}});
 
   typedef const boost::variant<string,vector<pair<int,string> > > combovar_t;
 
diff --git a/regression-tests.auth-py/test_LuaRecords.py b/regression-tests.auth-py/test_LuaRecords.py
index 973d98568c87..60082b0077d2 100644
--- a/regression-tests.auth-py/test_LuaRecords.py
+++ b/regression-tests.auth-py/test_LuaRecords.py
@@ -160,6 +160,11 @@ class TestLuaRecords(AuthTest):
 dblookup         IN    LUA    A  "dblookup('lookmeup.example.org', pdns.A)[1]"
 
 whitespace       IN    LUA    TXT "'foo" "bar'"
+
+geoipqueryattribute IN LUA    TXT ("string.format('%d %d %d %d %d %d %d',"
+                                   "GeoIPQueryAttribute.ASn, GeoIPQueryAttribute.City, GeoIPQueryAttribute.Continent,"
+                                   "GeoIPQueryAttribute.Country, GeoIPQueryAttribute.Country2, GeoIPQueryAttribute.Name,"
+                                   "GeoIPQueryAttribute.Region, GeoIPQueryAttribute.Location)")
         """,
         'createforward6.example.org': """
 createforward6.example.org.                 3600 IN SOA  {soa}
@@ -1109,6 +1114,22 @@ def testWhitespace(self, expectws=True):
         self.assertRcodeEqual(res, dns.rcode.NOERROR)
         self.assertEqual(res.answer, response.answer)
 
+    def testGeoIPQueryAttribute(self):
+        """
+        Test GeoIPQueryAttribute enum
+        """
+        name = 'geoipqueryattribute.example.org.'
+
+        query = dns.message.make_query(name, 'TXT')
+
+        response = dns.message.make_response(query)
+
+        response.answer.append(dns.rrset.from_text(name, 0, dns.rdataclass.IN, dns.rdatatype.TXT, '"0 1 2 3 4 5 6"'))
+
+        res = self.sendUDPQuery(query)
+        self.assertRcodeEqual(res, dns.rcode.NOERROR)
+        self.assertEqual(res.answer, response.answer)
+
 
 class TestLuaRecordsShared(TestLuaRecords):
     _config_template = """