From 71e0ad0bbe9bee1d33e75f54149c3bf33d3c5e55 Mon Sep 17 00:00:00 2001
From: Michael Tremer <michael.tremer@ipfire.org>
Date: Fri, 15 Nov 2019 14:18:59 +0000
Subject: [PATCH] location-query: Allow exporting data for nftables

Fixes: #12201
Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---
 man/location-query.txt       |  1 +
 src/python/location-query.in | 28 +++++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/man/location-query.txt b/man/location-query.txt
index 84f0fdc..ed6f4e5 100644
--- a/man/location-query.txt
+++ b/man/location-query.txt
@@ -50,6 +50,7 @@ or countries.
 	directly loaded into other software. The following options are available:
 	+
 	* 'list' (default): Just lists all networks, one per line
+	* 'nftables': For nftables
 	* 'xt_geoip': Returns a list of networks to be loaded into the xt_geoip
 	  kernel module
 
diff --git a/src/python/location-query.in b/src/python/location-query.in
index 258afbc..83f23ee 100644
--- 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:
-- 
2.39.2