src/python/location/export.py

   1 ###############################################################################
   2 #                                                                             #
   3 # libloc - A library to determine the location of someone on the Internet     #
   4 #                                                                             #
   5 # Copyright (C) 2020-2021 IPFire Development Team <info@ipfire.org>           #
   6 #                                                                             #
   7 # This library is free software; you can redistribute it and/or               #
   8 # modify it under the terms of the GNU Lesser General Public                  #
   9 # License as published by the Free Software Foundation; either                #
  10 # version 2.1 of the License, or (at your option) any later version.          #
  11 #                                                                             #
  12 # This library is distributed in the hope that it will be useful,             #
  13 # but WITHOUT ANY WARRANTY; without even the implied warranty of              #
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU           #
  15 # Lesser General Public License for more details.                             #
  16 #                                                                             #
  17 ###############################################################################
  18
  19 import io
  20 import ipaddress
  21 import logging
  22 import math
  23 import os
  24 import socket
  25 import sys
  26
  27 from .i18n import _
  28 import _location
  29
  30 # Initialise logging
  31 log = logging.getLogger("location.export")
  32 log.propagate = 1
  33
  34 FLAGS = {
  35         _location.NETWORK_FLAG_ANONYMOUS_PROXY    : "A1",
  36         _location.NETWORK_FLAG_SATELLITE_PROVIDER : "A2",
  37         _location.NETWORK_FLAG_ANYCAST            : "A3",
  38         _location.NETWORK_FLAG_DROP               : "XD",
  39 }
  40
  41 class OutputWriter(object):
  42         suffix = "networks"
  43         mode = "w"
  44
  45         def __init__(self, name, family=None, directory=None, f=None):
  46                 self.name = name
  47                 self.family = family
  48                 self.directory = directory
  49
  50                 # Tag
  51                 self.tag = self._make_tag()
  52
  53                 # Open output file
  54                 if f:
  55                         self.f = f
  56                 elif self.directory:
  57                         self.f = open(self.filename, self.mode)
  58                 elif "b" in self.mode:
  59                         self.f = io.BytesIO()
  60                 else:
  61                         self.f = io.StringIO()
  62
  63                 # Call any custom initialization
  64                 self.init()
  65
  66                 # Immediately write the header
  67                 self._write_header()
  68
  69         def init(self):
  70                 """
  71                         To be overwritten by anything that inherits from this
  72                 """
  73                 pass
  74
  75         def __repr__(self):
  76                 return "<%s %s f=%s>" % (self.__class__.__name__, self, self.f)
  77
  78         def _make_tag(self):
  79                 families = {
  80                         socket.AF_INET6 : "6",
  81                         socket.AF_INET  : "4",
  82                 }
  83
  84                 return "%sv%s" % (self.name, families.get(self.family, "?"))
  85
  86         @property
  87         def filename(self):
  88                 if self.directory:
  89                         return os.path.join(self.directory, "%s.%s" % (self.tag, self.suffix))
  90
  91         def _write_header(self):
  92                 """
  93                         The header of the file
  94                 """
  95                 pass
  96
  97         def _write_footer(self):
  98                 """
  99                         The footer of the file
 100                 """
 101                 pass
 102
 103         def write(self, network):
 104                 self.f.write("%s\n" % network)
 105
 106         def finish(self):
 107                 """
 108                         Called when all data has been written
 109                 """
 110                 self._write_footer()
 111
 112                 # Flush all output
 113                 self.f.flush()
 114
 115         def print(self):
 116                 """
 117                         Prints the entire output line by line
 118                 """
 119                 if isinstance(self.f, io.BytesIO):
 120                         raise TypeError(_("Won't write binary output to stdout"))
 121
 122                 # Go back to the beginning
 123                 self.f.seek(0)
 124
 125                 # Iterate over everything line by line
 126                 for line in self.f:
 127                         sys.stdout.write(line)
 128
 129
 130 class IpsetOutputWriter(OutputWriter):
 131         """
 132                 For ipset
 133         """
 134         suffix = "ipset"
 135
 136         # The value is being used if we don't know any better
 137         DEFAULT_HASHSIZE = 64
 138
 139         # We aim for this many networks in a bucket on average. This allows us to choose
 140         # how much memory we want to sacrifice to gain better performance. The lower the
 141         # factor, the faster a lookup will be, but it will use more memory.
 142         # We will aim for only using three quarters of all buckets to avoid any searches
 143         # through the linked lists.
 144         HASHSIZE_FACTOR = 0.75
 145
 146         def init(self):
 147                 # Count all networks
 148                 self.networks = 0
 149
 150                 # Check that family is being set
 151                 if not self.family:
 152                         raise ValueError("%s requires family being set" % self.__class__.__name__)
 153
 154         @property
 155         def hashsize(self):
 156                 """
 157                         Calculates an optimized hashsize
 158                 """
 159                 # Return the default value if we don't know the size of the set
 160                 if not self.networks:
 161                         return self.DEFAULT_HASHSIZE
 162
 163                 # Find the nearest power of two that is larger than the number of networks
 164                 # divided by the hashsize factor.
 165                 exponent = math.log(self.networks / self.HASHSIZE_FACTOR, 2)
 166
 167                 # Return the size of the hash (the minimum is 64)
 168                 return max(2 ** math.ceil(exponent), 64)
 169
 170         def _write_header(self):
 171                 # This must have a fixed size, because we will write the header again in the end
 172                 self.f.write("create %s hash:net family inet%s" % (
 173                         self.tag,
 174                         "6" if self.family == socket.AF_INET6 else ""
 175                 ))
 176                 self.f.write(" hashsize %8d maxelem 1048576 -exist\n" % self.hashsize)
 177                 self.f.write("flush %s\n" % self.tag)
 178
 179         def write(self, network):
 180                 self.f.write("add %s %s\n" % (self.tag, network))
 181
 182                 # Increment network counter
 183                 self.networks += 1
 184
 185         def _write_footer(self):
 186                 # Jump back to the beginning of the file
 187                 try:
 188                         self.f.seek(0)
 189
 190                 # If the output stream isn't seekable, we won't try writing the header again
 191                 except io.UnsupportedOperation:
 192                         return
 193
 194                 # Rewrite the header with better configuration
 195                 self._write_header()
 196
 197
 198 class NftablesOutputWriter(OutputWriter):
 199         """
 200                 For nftables
 201         """
 202         suffix = "set"
 203
 204         def _write_header(self):
 205                 self.f.write("define %s = {\n" % self.tag)
 206
 207         def _write_footer(self):
 208                 self.f.write("}\n")
 209
 210         def write(self, network):
 211                 self.f.write("  %s,\n" % network)
 212
 213
 214 class XTGeoIPOutputWriter(OutputWriter):
 215         """
 216                 Formats the output in that way, that it can be loaded by
 217                 the xt_geoip kernel module from xtables-addons.
 218         """
 219         mode = "wb"
 220
 221         @property
 222         def tag(self):
 223                 return self.name
 224
 225         @property
 226         def suffix(self):
 227                 return "iv%s" % ("6" if self.family == socket.AF_INET6 else "4")
 228
 229         def write(self, network):
 230                 self.f.write(network._first_address)
 231                 self.f.write(network._last_address)
 232
 233
 234 formats = {
 235         "ipset"    : IpsetOutputWriter,
 236         "list"     : OutputWriter,
 237         "nftables" : NftablesOutputWriter,
 238         "xt_geoip" : XTGeoIPOutputWriter,
 239 }
 240
 241 class Exporter(object):
 242         def __init__(self, db, writer):
 243                 self.db, self.writer = db, writer
 244
 245         def export(self, directory, families, countries, asns):
 246                 for family in families:
 247                         log.debug("Exporting family %s" % family)
 248
 249                         writers = {}
 250
 251                         # Create writers for countries
 252                         for country_code in countries:
 253                                 writers[country_code] = self.writer(country_code, family=family, directory=directory)
 254
 255                         # Create writers for ASNs
 256                         for asn in asns:
 257                                 writers[asn] = self.writer("AS%s" % asn, family=family, directory=directory)
 258
 259                         # Filter countries from special country codes
 260                         country_codes = [
 261                                 country_code for country_code in countries if not country_code in FLAGS.values()
 262                         ]
 263
 264                         # Get all networks that match the family
 265                         networks = self.db.search_networks(family=family,
 266                                 country_codes=country_codes, asns=asns, flatten=True)
 267
 268                         # Walk through all networks
 269                         for network in networks:
 270                                 # Write matching countries
 271                                 try:
 272                                         writers[network.country_code].write(network)
 273                                 except KeyError:
 274                                         pass
 275
 276                                 # Write matching ASNs
 277                                 try:
 278                                         writers[network.asn].write(network)
 279                                 except KeyError:
 280                                         pass
 281
 282                                 # Handle flags
 283                                 for flag in FLAGS:
 284                                         if network.has_flag(flag):
 285                                                 # Fetch the "fake" country code
 286                                                 country = FLAGS[flag]
 287
 288                                                 try:
 289                                                         writers[country].write(network)
 290                                                 except KeyError:
 291                                                         pass
 292
 293                         # Write everything to the filesystem
 294                         for writer in writers.values():
 295                                 writer.finish()
 296
 297                         # Print to stdout
 298                         if not directory:
 299                                 for writer in writers.values():
 300                                         writer.print()