location-query: Allow exporting data for nftables

[location/libloc.git] / src / python / location-query.in

diff --git a/src/python/location-query.in b/src/python/location-query.in
index 258afbc8f878bf23713af64cfc5c8f501373f260..83f23ee1209a0e243a4609ee747e3f705d257d65 100644 (file)
--- a/src/python/location-query.in
+++ b/src/python/location-query.in
@@ -38,6 +38,9 @@ def _(singular, plural=None, n=None):
 # Output formatters
 
 class OutputFormatter(object):
+       def __init__(self, ns):
+               self.ns = ns
+
        def __enter__(self):
                # Open the output
                self.open()
@@ -48,6 +51,14 @@ class OutputFormatter(object):
                if tb is None:
                        self.close()
 
+       @property
+       def name(self):
+               if "country_code" in self.ns:
+                       return "networks_country_%s" % self.ns.country_code[0]
+
+               elif "asn" in self.ns:
+                       return "networks_AS%s" % self.ns.asn[0]
+
        def open(self):
                pass
 
@@ -58,6 +69,20 @@ class OutputFormatter(object):
                print(network)
 
 
+class NftablesOutputFormatter(OutputFormatter):
+       """
+               For nftables
+       """
+       def open(self):
+               print("define %s = {" % self.name)
+
+       def close(self):
+               print("}")
+
+       def network(self, network):
+               print(" %s," % network)
+
+
 class XTGeoIPOutputFormatter(OutputFormatter):
        """
                Formats the output in that way, that it can be loaded by
@@ -78,6 +103,7 @@ class XTGeoIPOutputFormatter(OutputFormatter):
 class CLI(object):
        output_formats = {
                "list"     : OutputFormatter,
+               "nftables" : NftablesOutputFormatter,
                "xt_geoip" : XTGeoIPOutputFormatter,
        }
 
@@ -246,7 +272,7 @@ class CLI(object):
                except KeyError:
                        cls = OutputFormatter
 
-               return cls()
+               return cls(ns)
 
        def handle_list_networks_by_as(self, db, ns):
                with self.__get_output_formatter(ns) as f: