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 import _location
  27
  28 # Initialise logging
  29 log = logging.getLogger("location.export")
  30 log.propagate = 1
  31
  32 FLAGS = {
  33         _location.NETWORK_FLAG_ANONYMOUS_PROXY    : "A1",
  34         _location.NETWORK_FLAG_SATELLITE_PROVIDER : "A2",
  35         _location.NETWORK_FLAG_ANYCAST            : "A3",
  36 }
  37
  38 class OutputWriter(object):
  39         suffix = "networks"
  40         mode = "w"
  41
  42         def __init__(self, db, f, prefix=None, flatten=True):
  43                 self.db, self.f, self.prefix, self.flatten = db, f, prefix, flatten
  44
  45                 # The previously written network
  46                 self._last_network = None
  47
  48                 # Immediately write the header
  49                 self._write_header()
  50
  51         @classmethod
  52         def open(cls, db, filename, **kwargs):
  53                 """
  54                         Convenience function to open a file
  55                 """
  56                 f = open(filename, cls.mode)
  57
  58                 return cls(db, f, **kwargs)
  59
  60         def __repr__(self):
  61                 return "<%s f=%s>" % (self.__class__.__name__, self.f)
  62
  63         def _flatten(self, network):
  64                 """
  65                         Checks if the given network needs to be written to file,
  66                         or if it is a subnet of the previously written network.
  67                 """
  68                 if self._last_network and network.is_subnet_of(self._last_network):
  69                         return True
  70
  71                 # Remember this network for the next call
  72                 self._last_network = network
  73                 return False
  74
  75         def _write_header(self):
  76                 """
  77                         The header of the file
  78                 """
  79                 pass
  80
  81         def _write_footer(self):
  82                 """
  83                         The footer of the file
  84                 """
  85                 pass
  86
  87         def _write_network(self, network):
  88                 self.f.write("%s\n" % network)
  89
  90         def write(self, network):
  91                 if self.flatten and self._flatten(network):
  92                         log.debug("Skipping writing network %s (last one was %s)" % (network, self._last_network))
  93                         return
  94
  95                 return self._write_network(network)
  96
  97         def finish(self):
  98                 """
  99                         Called when all data has been written
 100                 """
 101                 self._write_footer()
 102
 103                 # Close the file
 104                 self.f.close()
 105
 106
 107 class IpsetOutputWriter(OutputWriter):
 108         """
 109                 For ipset
 110         """
 111         suffix = "ipset"
 112
 113         def _write_header(self):
 114                 self.f.write("create %s hash:net family inet hashsize 1024 maxelem 65536\n" % self.prefix)
 115
 116         def _write_network(self, network):
 117                 self.f.write("add %s %s\n" % (self.prefix, network))
 118
 119
 120 class NftablesOutputWriter(OutputWriter):
 121         """
 122                 For nftables
 123         """
 124         suffix = "set"
 125
 126         def _write_header(self):
 127                 self.f.write("define %s = {\n" % self.prefix)
 128
 129         def _write_footer(self):
 130                 self.f.write("}\n")
 131
 132         def _write_network(self, network):
 133                 self.f.write("  %s,\n" % network)
 134
 135
 136 class XTGeoIPOutputWriter(OutputWriter):
 137         """
 138                 Formats the output in that way, that it can be loaded by
 139                 the xt_geoip kernel module from xtables-addons.
 140         """
 141         suffix = "iv"
 142         mode = "wb"
 143
 144         def _write_network(self, network):
 145                 for address in (network.first_address, network.last_address):
 146                         # Convert this into a string of bits
 147                         bytes = socket.inet_pton(
 148                                 socket.AF_INET6 if network.version == 6 else socket.AF_INET, "%s" % address,
 149                         )
 150
 151                         self.f.write(bytes)
 152
 153
 154 formats = {
 155         "ipset"    : IpsetOutputWriter,
 156         "list"     : OutputWriter,
 157         "nftables" : NftablesOutputWriter,
 158         "xt_geoip" : XTGeoIPOutputWriter,
 159 }
 160
 161 class Exporter(object):
 162         def __init__(self, db, writer):
 163                 self.db, self.writer = db, writer
 164
 165         def export(self, directory, families, countries, asns):
 166                 for family in families:
 167                         log.debug("Exporting family %s" % family)
 168
 169                         writers = {}
 170
 171                         # Create writers for countries
 172                         for country_code in countries:
 173                                 filename = self._make_filename(
 174                                         directory, prefix=country_code, suffix=self.writer.suffix, family=family,
 175                                 )
 176
 177                                 writers[country_code] = self.writer.open(self.db, filename, prefix="CC_%s" % country_code)
 178
 179                         # Create writers for ASNs
 180                         for asn in asns:
 181                                 filename = self._make_filename(
 182                                         directory, "AS%s" % asn, suffix=self.writer.suffix, family=family,
 183                                 )
 184
 185                                 writers[asn] = self.writer.open(self.db, filename, prefix="AS%s" % asn)
 186
 187                         # Filter countries from special country codes
 188                         country_codes = [
 189                                 country_code for country_code in countries if not country_code in FLAGS.values()
 190                         ]
 191
 192                         # Collect flags
 193                         flags = 0
 194                         for flag in FLAGS:
 195                                 if FLAGS[flag] in countries:
 196                                         flags |= flag
 197
 198                         # Get all networks that match the family
 199                         networks = self.db.search_networks(family=family,
 200                                 country_codes=country_codes, flags=flags, flatten=True)
 201
 202                         # Walk through all networks
 203                         for network in networks:
 204                                 # Write matching countries
 205                                 try:
 206                                         writers[network.country_code].write(network)
 207                                 except KeyError:
 208                                         pass
 209
 210                                 # Write matching ASNs
 211                                 try:
 212                                         writers[network.asn].write(network)
 213                                 except KeyError:
 214                                         pass
 215
 216                                 # Handle flags
 217                                 for flag in FLAGS:
 218                                         if network.has_flag(flag):
 219                                                 # Fetch the "fake" country code
 220                                                 country = FLAGS[flag]
 221
 222                                                 try:
 223                                                         writers[country].write(network)
 224                                                 except KeyError:
 225                                                         pass
 226
 227                         # Write everything to the filesystem
 228                         for writer in writers.values():
 229                                 writer.finish()
 230
 231         def _make_filename(self, directory, prefix, suffix, family):
 232                 filename = "%s.%s%s" % (
 233                         prefix, suffix, "6" if family == socket.AF_INET6 else "4"
 234                 )
 235
 236                 return os.path.join(directory, filename)