src/python/export.py

   1 #!/usr/bin/python3
   2 ###############################################################################
   3 #                                                                             #
   4 # libloc - A library to determine the location of someone on the Internet     #
   5 #                                                                             #
   6 # Copyright (C) 2020 IPFire Development Team <info@ipfire.org>                #
   7 #                                                                             #
   8 # This library is free software; you can redistribute it and/or               #
   9 # modify it under the terms of the GNU Lesser General Public                  #
  10 # License as published by the Free Software Foundation; either                #
  11 # version 2.1 of the License, or (at your option) any later version.          #
  12 #                                                                             #
  13 # This library is distributed in the hope that it will be useful,             #
  14 # but WITHOUT ANY WARRANTY; without even the implied warranty of              #
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU           #
  16 # Lesser General Public License for more details.                             #
  17 #                                                                             #
  18 ###############################################################################
  19
  20 import io
  21 import ipaddress
  22 import logging
  23 import os
  24 import socket
  25
  26 # Initialise logging
  27 log = logging.getLogger("location.export")
  28 log.propagate = 1
  29
  30 class OutputWriter(object):
  31         suffix = "networks"
  32         mode = "w"
  33
  34         def __init__(self, f, prefix=None):
  35                 self.f, self.prefix = f, prefix
  36
  37                 # Immediately write the header
  38                 self._write_header()
  39
  40         @classmethod
  41         def open(cls, filename, **kwargs):
  42                 """
  43                         Convenience function to open a file
  44                 """
  45                 f = open(filename, cls.mode)
  46
  47                 return cls(f, **kwargs)
  48
  49         def __repr__(self):
  50                 return "<%s f=%s>" % (self.__class__.__name__, self.f)
  51
  52         def _write_header(self):
  53                 """
  54                         The header of the file
  55                 """
  56                 pass
  57
  58         def _write_footer(self):
  59                 """
  60                         The footer of the file
  61                 """
  62                 pass
  63
  64         def write(self, network):
  65                 self.f.write("%s\n" % network)
  66
  67         def finish(self):
  68                 """
  69                         Called when all data has been written
  70                 """
  71                 self._write_footer()
  72
  73                 # Close the file
  74                 self.f.close()
  75
  76
  77 class IpsetOutputWriter(OutputWriter):
  78         """
  79                 For ipset
  80         """
  81         suffix = "ipset"
  82
  83         def _write_header(self):
  84                 self.f.write("create %s hash:net family inet hashsize 1024 maxelem 65536\n" % self.prefix)
  85
  86         def write(self, network):
  87                 self.f.write("add %s %s\n" % (self.prefix, network))
  88
  89
  90 class NftablesOutputWriter(OutputWriter):
  91         """
  92                 For nftables
  93         """
  94         suffix = "set"
  95
  96         def _write_header(self):
  97                 self.f.write("define %s = {\n" % self.prefix)
  98
  99         def _write_footer(self):
 100                 self.f.write("}\n")
 101
 102         def write(self, network):
 103                 self.f.write("  %s,\n" % network)
 104
 105
 106 class XTGeoIPOutputWriter(OutputWriter):
 107         """
 108                 Formats the output in that way, that it can be loaded by
 109                 the xt_geoip kernel module from xtables-addons.
 110         """
 111         suffix = "iv"
 112         mode = "wb"
 113
 114         def write(self, network):
 115                 n = ipaddress.ip_network("%s" % network)
 116
 117                 for address in (n.network_address, n.broadcast_address):
 118                         bytes = socket.inet_pton(
 119                                 socket.AF_INET6 if address.version == 6 else socket.AF_INET,
 120                                 "%s" % address,
 121                         )
 122
 123                         self.f.write(bytes)
 124
 125
 126 formats = {
 127         "ipset"    : IpsetOutputWriter,
 128         "list"     : OutputWriter,
 129         "nftables" : NftablesOutputWriter,
 130         "xt_geoip" : XTGeoIPOutputWriter,
 131 }
 132
 133 class Exporter(object):
 134         def __init__(self, db, writer):
 135                 self.db, self.writer = db, writer
 136
 137         def export(self, directory, families, countries, asns):
 138                 for family in families:
 139                         log.debug("Exporting family %s" % family)
 140
 141                         writers = {}
 142
 143                         # Create writers for countries
 144                         for country_code in countries:
 145                                 filename = self._make_filename(
 146                                         directory, prefix=country_code, suffix=self.writer.suffix, family=family,
 147                                 )
 148
 149                                 writers[country_code] = self.writer.open(filename, prefix="CC_%s" % country_code)
 150
 151                         # Create writers for ASNs
 152                         for asn in asns:
 153                                 filename = self._make_filename(
 154                                         directory, "AS%s" % asn, suffix=self.writer.suffix, family=family,
 155                                 )
 156
 157                                 writers[asn] = self.writer.open(filename, prefix="AS%s" % asn)
 158
 159                         # Get all networks that match the family
 160                         networks = self.db.search_networks(family=family)
 161
 162                         # Walk through all networks
 163                         for network in networks:
 164                                 # Write matching countries
 165                                 try:
 166                                         writers[network.country_code].write(network)
 167                                 except KeyError:
 168                                         pass
 169
 170                                 # Write matching ASNs
 171                                 try:
 172                                         writers[network.asn].write(network)
 173                                 except KeyError:
 174                                         pass
 175
 176                         # Write everything to the filesystem
 177                         for writer in writers.values():
 178                                 writer.finish()
 179
 180         def _make_filename(self, directory, prefix, suffix, family):
 181                 filename = "%s.%s%s" % (
 182                         prefix, suffix, "6" if family == socket.AF_INET6 else "4"
 183                 )
 184
 185                 return os.path.join(directory, filename)