From f94234194abfdb9201430c767d7372282f4a056c Mon Sep 17 00:00:00 2001
From: Peter van Dijk <peter.van.dijk@powerdns.com>
Date: Fri, 16 Mar 2018 13:56:54 +0100
Subject: [PATCH] hook LUA to geoip again

---
 modules/geoipbackend/geoipbackend.cc   | 37 ++++++++++++++++++++++++--
 modules/geoipbackend/geoipinterface.hh |  5 +++-
 pdns/lua-record.cc                     | 14 +++++-----
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/modules/geoipbackend/geoipbackend.cc b/modules/geoipbackend/geoipbackend.cc
index 58d662099e..6962b39be9 100644
--- a/modules/geoipbackend/geoipbackend.cc
+++ b/modules/geoipbackend/geoipbackend.cc
@@ -50,7 +50,7 @@ public:
 };
 
 static vector<GeoIPDomain> s_domains;
-static int s_rc = 0; // refcount
+static int s_rc = 0; // refcount - always accessed under lock
 
 static string GeoIP_WEEKDAYS[] = { "mon", "tue", "wed", "thu", "fri", "sat", "sun" };
 static string GeoIP_MONTHS[] = { "jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec" };
@@ -84,6 +84,8 @@ GeoIPBackend::GeoIPBackend(const string& suffix) {
 
 static vector<std::unique_ptr<GeoIPInterface> > s_geoip_files;
 
+string getGeoForLua(const std::string& ip, int qaint);
+
 void GeoIPBackend::initialize() {
   YAML::Node config;
   vector<GeoIPDomain> tmp_domains;
@@ -101,7 +103,8 @@ void GeoIPBackend::initialize() {
   if (s_geoip_files.empty())
     L<<Logger::Warning<<"No GeoIP database files loaded!"<<endl;
 
-  config = YAML::LoadFile(getArg("zones-file"));
+  if(!getArg("zones-file").empty())
+    config = YAML::LoadFile(getArg("zones-file"));
 
   for(YAML::Node domain :  config["domains"]) {
     GeoIPDomain dom;
@@ -264,6 +267,9 @@ void GeoIPBackend::initialize() {
 
   s_domains.clear();
   std::swap(s_domains, tmp_domains);
+
+  extern std::function<std::string(const std::string& ip, int)> g_getGeo;
+  g_getGeo = getGeoForLua;
 }
 
 GeoIPBackend::~GeoIPBackend() {
@@ -434,6 +440,13 @@ string queryGeoIP(const string &ip, bool v6, GeoIPInterface::GeoIPQueryAttribute
       if (v6) found = gi->queryCityV6(val, gl, ip);
       else found = gi->queryCity(val, gl, ip);
       break;
+    case GeoIPInterface::Location:
+      double lat=0, lon=0;
+      boost::optional<int> alt, prec;
+      if (v6) found = gi->queryLocationV6(gl, ip, lat, lon, alt, prec);
+      else found = gi->queryLocation(gl, ip, lat, lon, alt, prec);
+      val = std::to_string(lat)+" "+std::to_string(lon);
+      break;
     }
 
     if (!found || val.empty() || val == "--") continue; // try next database
@@ -446,6 +459,26 @@ string queryGeoIP(const string &ip, bool v6, GeoIPInterface::GeoIPQueryAttribute
   return ret;
 }
 
+string getGeoForLua(const std::string& ip, int qaint)
+{
+  GeoIPInterface::GeoIPQueryAttribute qa((GeoIPInterface::GeoIPQueryAttribute)qaint);
+  try {
+    GeoIPNetmask gl;
+    string res=queryGeoIP(ip, false, qa, gl);
+    //    cout<<"Result for "<<ip<<" lookup: "<<res<<endl;
+    if(qa==GeoIPInterface::ASn && boost::starts_with(res, "as"))
+      return res.substr(2);
+    return res;
+  }
+  catch(std::exception& e) {
+    cout<<"Error: "<<e.what()<<endl;
+  }
+  catch(PDNSException& e) {
+    cout<<"Error: "<<e.reason<<endl;
+  }
+  return "";
+}
+
 bool queryGeoLocation(const string &ip, bool v6, GeoIPNetmask& gl, double& lat, double& lon,
                       boost::optional<int>& alt, boost::optional<int>& prec)
 {
diff --git a/modules/geoipbackend/geoipinterface.hh b/modules/geoipbackend/geoipinterface.hh
index 87d83bf4fb..7b1e1805a2 100644
--- a/modules/geoipbackend/geoipinterface.hh
+++ b/modules/geoipbackend/geoipinterface.hh
@@ -24,6 +24,8 @@
 
 #include "boost/optional.hpp"
 
+#include "geoipbackend.hh"
+
 class GeoIPInterface {
 public:
   enum GeoIPQueryAttribute {
@@ -33,7 +35,8 @@ public:
     Country,
     Country2,
     Name,
-    Region
+    Region,
+    Location
   };
 
   virtual bool queryCountry(string &ret, GeoIPNetmask& gl, const string &ip) = 0;
diff --git a/pdns/lua-record.cc b/pdns/lua-record.cc
index e5f9c6ee06..8f21f83d9f 100644
--- a/pdns/lua-record.cc
+++ b/pdns/lua-record.cc
@@ -7,7 +7,7 @@
 #include "ueberbackend.hh"
 #include <boost/format.hpp>
 
-#include "../modules/geoipbackend/geoipbackend.hh" // only for the enum
+#include "../modules/geoipbackend/geoipinterface.hh" // only for the enum
 
 /* to do:
    block AXFR unless TSIG, or override
@@ -232,7 +232,7 @@ bool doCompare(const T& var, const std::string& res, const C& cmp)
 }
 
 
-std::string getGeo(const std::string& ip, GeoIPBackend::GeoIPQueryAttribute qa)
+std::string getGeo(const std::string& ip, GeoIPInterface::GeoIPQueryAttribute qa)
 {
   static bool initialized;
   extern std::function<std::string(const std::string& ip, int)> g_getGeo;
@@ -288,7 +288,7 @@ static ComboAddress whashed(const ComboAddress& bestwho, vector<pair<int,ComboAd
 
 static bool getLatLon(const std::string& ip, double& lat, double& lon)
 {
-  string inp = getGeo(ip, GeoIPBackend::LatLon);
+  string inp = getGeo(ip, GeoIPInterface::Location);
   if(inp.empty())
     return false;
   lat=atof(inp.c_str());
@@ -735,27 +735,27 @@ std::vector<shared_ptr<DNSRecordContent>> luaSynth(const std::string& code, cons
   lua.executeCode("debug.sethook(report, '', 1000)");
 
   // TODO: make this better. Accept netmask/CA objects; provide names for the attr constants
-  lua.writeFunction("geoiplookup", [](const string &ip, const GeoIPBackend::GeoIPQueryAttribute attr) {
+  lua.writeFunction("geoiplookup", [](const string &ip, const GeoIPInterface::GeoIPQueryAttribute attr) {
     return getGeo(ip, attr);
   });
 
   typedef const boost::variant<string,vector<pair<int,string> > > combovar_t;
   lua.writeFunction("continent", [&bestwho](const combovar_t& continent) {
-      string res=getGeo(bestwho.toString(), GeoIPBackend::Continent);
+      string res=getGeo(bestwho.toString(), GeoIPInterface::Continent);
       return doCompare(continent, res, [](const std::string& a, const std::string& b) {
           return !strcasecmp(a.c_str(), b.c_str());
         });
     });
 
   lua.writeFunction("asnum", [&bestwho](const combovar_t& asns) {
-      string res=getGeo(bestwho.toString(), GeoIPBackend::ASn);
+      string res=getGeo(bestwho.toString(), GeoIPInterface::ASn);
       return doCompare(asns, res, [](const std::string& a, const std::string& b) {
           return !strcasecmp(a.c_str(), b.c_str());
         });
     });
   
   lua.writeFunction("country", [&bestwho](const combovar_t& var) {
-      string res = getGeo(bestwho.toString(), GeoIPBackend::Country2);
+      string res = getGeo(bestwho.toString(), GeoIPInterface::Country2);
       return doCompare(var, res, [](const std::string& a, const std::string& b) {
           return !strcasecmp(a.c_str(), b.c_str());
         });
-- 
2.47.2