From: Michael Tremer <michael.tremer@ipfire.org>
Date: Fri, 15 Nov 2019 13:53:02 +0000 (+0000)
Subject: location-query: Support listing networks for xt_geoip
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=4439e3173767e72435c1ef8aa59a24b9b789d7f3;p=people%2Fsennis%2Flibloc.git

location-query: Support listing networks for xt_geoip

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---

diff --git a/man/location-query.txt b/man/location-query.txt
index 216d1db..84f0fdc 100644
--- a/man/location-query.txt
+++ b/man/location-query.txt
@@ -43,13 +43,22 @@ or countries.
 	+
 	The search will be performed case-insensitively.
 
-'list-networks-by-as ASN'::
+'list-networks-by-as [--output-format FORMAT] ASN'::
 	Lists all networks which belong to this Autonomous System.
+	+
+	The '--output-format' parameter can change the output so that it can be
+	directly loaded into other software. The following options are available:
+	+
+	* 'list' (default): Just lists all networks, one per line
+	* 'xt_geoip': Returns a list of networks to be loaded into the xt_geoip
+	  kernel module
 
-'list-networks-by-cc COUNTRY_CODE'::
+'list-networks-by-cc [--output-format FORMAT] COUNTRY_CODE'::
 	Lists all networks that belong to a country.
 	+
 	The country has to be encoded in ISO3166 Alpha-2 notation.
+	+
+	See above for usage of the '--output-format' parameter.
 
 '--help'::
 	Shows a short help text on using this program.
diff --git a/src/python/location-query.in b/src/python/location-query.in
index 2172da2..258afbc 100644
--- a/src/python/location-query.in
+++ b/src/python/location-query.in
@@ -19,6 +19,9 @@
 
 import argparse
 import gettext
+import ipaddress
+import os
+import socket
 import sys
 import syslog
 
@@ -32,7 +35,52 @@ def _(singular, plural=None, n=None):
 
 	return gettext.dgettext("libloc", singular)
 
+# Output formatters
+
+class OutputFormatter(object):
+	def __enter__(self):
+		# Open the output
+		self.open()
+
+		return self
+
+	def __exit__(self, type, value, tb):
+		if tb is None:
+			self.close()
+
+	def open(self):
+		pass
+
+	def close(self):
+		pass
+
+	def network(self, network):
+		print(network)
+
+
+class XTGeoIPOutputFormatter(OutputFormatter):
+	"""
+		Formats the output in that way, that it can be loaded by
+		the xt_geoip kernel module from xtables-addons.
+	"""
+	def network(self, network):
+		n = ipaddress.ip_network("%s" % network)
+
+		for address in (n.network_address, n.broadcast_address):
+			bytes = socket.inet_pton(
+				socket.AF_INET6 if address.version == 6 else socket.AF_INET,
+				"%s" % address,
+			)
+
+			os.write(1, bytes)
+
+
 class CLI(object):
+	output_formats = {
+		"list"     : OutputFormatter,
+		"xt_geoip" : XTGeoIPOutputFormatter,
+	}
+
 	def parse_cli(self):
 		parser = argparse.ArgumentParser(
 			description=_("Location Database Command Line Interface"),
@@ -78,6 +126,8 @@ class CLI(object):
 			help=_("Lists all networks in an AS"),
 		)
 		list_networks_by_as.add_argument("asn", nargs=1, type=int)
+		list_networks_by_as.add_argument("--output-format",
+			choices=self.output_formats.keys(), default="list")
 		list_networks_by_as.set_defaults(func=self.handle_list_networks_by_as)
 
 		# List all networks in a country
@@ -85,6 +135,8 @@ class CLI(object):
 			help=_("Lists all networks in a country"),
 		)
 		list_networks_by_cc.add_argument("country_code", nargs=1)
+		list_networks_by_cc.add_argument("--output-format",
+			choices=self.output_formats.keys(), default="list")
 		list_networks_by_cc.set_defaults(func=self.handle_list_networks_by_cc)
 
 		args = parser.parse_args()
@@ -188,17 +240,28 @@ class CLI(object):
 			for a in db.search_as(query):
 				print(a)
 
+	def __get_output_formatter(self, ns):
+		try:
+			cls = self.output_formats[ns.output_format]
+		except KeyError:
+			cls = OutputFormatter
+
+		return cls()
+
 	def handle_list_networks_by_as(self, db, ns):
-		for asn in ns.asn:
-			# Print all matching networks
-			for n in db.search_networks(asn=asn):
-				print(n)
+		with self.__get_output_formatter(ns) as f:
+			for asn in ns.asn:
+				# Print all matching networks
+				for n in db.search_networks(asn=asn):
+					f.network(n)
 
 	def handle_list_networks_by_cc(self, db, ns):
-		for country_code in ns.country_code:
-			# Print all matching networks
-			for n in db.search_networks(country_code=country_code):
-				print(n)
+		with self.__get_output_formatter(ns) as f:
+			for country_code in ns.country_code:
+				# Print all matching networks
+				for n in db.search_networks(country_code=country_code):
+					f.network(n)
+
 
 def main():
 	# Run the command line interface