Add GeoIPQueryAttribute to LUA

BozhanL · Dec 13, 2024 · 0b45ac7 · 0b45ac7
1 parent 5089e2c
commit 0b45ac7
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git 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
@@ -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 = """