]>
Commit | Line | Data |
---|---|---|
6838b55a MT |
1 | #!/usr/bin/python3 |
2 | ############################################################################### | |
3 | # # | |
4 | # libloc - A library to determine the location of someone on the Internet # | |
5 | # # | |
6 | # Copyright (C) 2019 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 argparse | |
6838b55a MT |
21 | import io |
22 | import ipaddress | |
23 | import logging | |
6838b55a | 24 | import os.path |
141b1099 | 25 | import re |
6838b55a MT |
26 | import socket |
27 | import sys | |
28 | ||
29 | # Load our location module | |
30 | import location | |
7dccb767 | 31 | from location.i18n import _ |
6838b55a | 32 | |
6838b55a | 33 | # Initialise logging |
e44b30f4 MT |
34 | log = logging.getLogger("location.exporter") |
35 | log.propagate = 1 | |
6838b55a | 36 | |
6838b55a MT |
37 | class OutputWriter(object): |
38 | suffix = "networks" | |
39 | ||
40 | def __init__(self, family, country_code=None, asn=None): | |
41 | self.family, self.country_code, self.asn = family, country_code, asn | |
42 | ||
43 | self.f = io.BytesIO() | |
44 | ||
45 | def write_out(self, directory): | |
46 | # Make the output filename | |
47 | filename = os.path.join( | |
48 | directory, self._make_filename(), | |
49 | ) | |
50 | ||
51 | with open(filename, "wb") as f: | |
52 | self._write_header(f) | |
53 | ||
54 | # Copy all data into the file | |
55 | f.write(self.f.getbuffer()) | |
56 | ||
57 | self._write_footer(f) | |
58 | ||
59 | def _make_filename(self): | |
60 | return "%s.%s%s" % ( | |
61 | self.country_code or "AS%s" % self.asn, | |
62 | self.suffix, | |
63 | "6" if self.family == socket.AF_INET6 else "4" | |
64 | ) | |
65 | ||
66 | @property | |
67 | def name(self): | |
68 | if self.country_code: | |
69 | return "CC_%s" % self.country_code | |
70 | ||
71 | if self.asn: | |
72 | return "AS%s" % self.asn | |
73 | ||
74 | def _write_header(self, f): | |
75 | """ | |
76 | The header of the file | |
77 | """ | |
78 | pass | |
79 | ||
80 | def _write_footer(self, f): | |
81 | """ | |
82 | The footer of the file | |
83 | """ | |
84 | pass | |
85 | ||
86 | def write(self, network): | |
87 | s = "%s\n" % network | |
88 | ||
89 | self.f.write(s.encode("ascii")) | |
90 | ||
91 | ||
92 | class IpsetOutputWriter(OutputWriter): | |
93 | """ | |
94 | For ipset | |
95 | """ | |
96 | suffix = "ipset" | |
97 | ||
98 | def _write_header(self, f): | |
a4ba4762 | 99 | h = "create %s hash:net family inet hashsize 1024 maxelem 65536\n" % self.name |
6838b55a MT |
100 | |
101 | f.write(h.encode("ascii")) | |
102 | ||
103 | def write(self, network): | |
104 | s = "add %s %s\n" % (self.name, network) | |
105 | ||
106 | self.f.write(s.encode("ascii")) | |
107 | ||
108 | ||
109 | class NftablesOutputWriter(OutputWriter): | |
110 | """ | |
111 | For nftables | |
112 | """ | |
113 | suffix = "set" | |
114 | ||
115 | def _write_header(self, f): | |
a28d854c | 116 | h = "define %s = {\n" % self.name |
6838b55a MT |
117 | |
118 | f.write(h.encode("ascii")) | |
119 | ||
a28d854c | 120 | def _write_footer(self, f): |
6838b55a MT |
121 | f.write(b"}") |
122 | ||
123 | def write(self, network): | |
124 | s = " %s,\n" % network | |
125 | ||
126 | self.f.write(s.encode("ascii")) | |
127 | ||
128 | ||
129 | class XTGeoIPOutputWriter(OutputWriter): | |
130 | """ | |
131 | Formats the output in that way, that it can be loaded by | |
132 | the xt_geoip kernel module from xtables-addons. | |
133 | """ | |
134 | suffix = "iv" | |
135 | ||
136 | def write(self, network): | |
137 | n = ipaddress.ip_network("%s" % network) | |
138 | ||
139 | for address in (n.network_address, n.broadcast_address): | |
140 | bytes = socket.inet_pton( | |
141 | socket.AF_INET6 if address.version == 6 else socket.AF_INET, | |
142 | "%s" % address, | |
143 | ) | |
144 | ||
145 | self.f.write(bytes) | |
146 | ||
147 | ||
148 | class Exporter(object): | |
149 | def __init__(self, db, writer): | |
150 | self.db = db | |
151 | self.writer = writer | |
152 | ||
48bcc19d SS |
153 | def export(self, directory, families, countries, asns): |
154 | for family in families: | |
6838b55a MT |
155 | log.debug("Exporting family %s" % family) |
156 | ||
157 | writers = {} | |
158 | ||
159 | # Create writers for countries | |
160 | for country_code in countries: | |
161 | writers[country_code] = self.writer(family, country_code=country_code) | |
162 | ||
163 | # Create writers for ASNs | |
164 | for asn in asns: | |
165 | writers[asn] = self.writer(family, asn=asn) | |
166 | ||
167 | # Get all networks that match the family | |
168 | networks = self.db.search_networks(family=family) | |
169 | ||
170 | # Walk through all networks | |
171 | for network in networks: | |
172 | # Write matching countries | |
173 | if network.country_code in countries: | |
174 | writers[network.country_code].write(network) | |
175 | ||
176 | # Write matching ASNs | |
177 | if network.asn in asns: | |
178 | writers[network.asn].write(network) | |
179 | ||
180 | # Write everything to the filesystem | |
181 | for writer in writers.values(): | |
182 | writer.write_out(directory) | |
183 | ||
184 | ||
185 | class CLI(object): | |
186 | output_formats = { | |
187 | "ipset" : IpsetOutputWriter, | |
188 | "list" : OutputWriter, | |
189 | "nftables" : NftablesOutputWriter, | |
190 | "xt_geoip" : XTGeoIPOutputWriter, | |
191 | } | |
192 | ||
193 | def parse_cli(self): | |
194 | parser = argparse.ArgumentParser( | |
195 | description=_("Location Exporter Command Line Interface"), | |
196 | ) | |
197 | ||
198 | # Global configuration flags | |
199 | parser.add_argument("--debug", action="store_true", | |
200 | help=_("Enable debug output")) | |
bc1f5f53 MT |
201 | parser.add_argument("--quiet", action="store_true", |
202 | help=_("Enable quiet mode")) | |
6838b55a MT |
203 | |
204 | # version | |
205 | parser.add_argument("--version", action="version", | |
d2714e4a | 206 | version="%(prog)s @VERSION@") |
6838b55a MT |
207 | |
208 | # database | |
209 | parser.add_argument("--database", "-d", | |
210 | default="@databasedir@/database.db", help=_("Path to database"), | |
211 | ) | |
212 | ||
213 | # format | |
214 | parser.add_argument("--format", help=_("Output format"), | |
8fe4a24a | 215 | default="list", choices=self.output_formats.keys()) |
6838b55a MT |
216 | |
217 | # directory | |
218 | parser.add_argument("--directory", help=_("Output directory"), required=True) | |
219 | ||
48bcc19d SS |
220 | # family |
221 | parser.add_argument("--family", help=_("Specify address family"), choices=("ipv6", "ipv4")) | |
222 | ||
6838b55a MT |
223 | # Countries and Autonomous Systems |
224 | parser.add_argument("objects", nargs="+") | |
225 | ||
226 | args = parser.parse_args() | |
227 | ||
bc1f5f53 | 228 | # Configure logging |
6838b55a | 229 | if args.debug: |
f9de5e61 | 230 | location.logger.set_level(logging.DEBUG) |
bc1f5f53 MT |
231 | elif args.quiet: |
232 | location.logger.set_level(logging.WARNING) | |
6838b55a MT |
233 | |
234 | return args | |
235 | ||
236 | def run(self): | |
237 | # Parse command line arguments | |
238 | args = self.parse_cli() | |
239 | ||
240 | # Call function | |
241 | ret = self.handle_export(args) | |
242 | ||
243 | # Return with exit code | |
244 | if ret: | |
245 | sys.exit(ret) | |
246 | ||
247 | # Otherwise just exit | |
248 | sys.exit(0) | |
249 | ||
250 | def handle_export(self, ns): | |
251 | countries, asns = [], [] | |
252 | ||
48bcc19d SS |
253 | # Translate family |
254 | if ns.family == "ipv6": | |
255 | families = [ socket.AF_INET6 ] | |
256 | elif ns.family == "ipv4": | |
257 | families = [ socket.AF_INET ] | |
258 | else: | |
259 | families = [ socket.AF_INET6, socket.AF_INET ] | |
260 | ||
6838b55a | 261 | for object in ns.objects: |
141b1099 MT |
262 | m = re.match("^AS(\d+)$", object) |
263 | if m: | |
264 | object = int(m.group(1)) | |
6838b55a MT |
265 | |
266 | asns.append(object) | |
267 | ||
8301f971 MT |
268 | elif location.country_code_is_valid(object) \ |
269 | or object in ("A1", "A2", "A3"): | |
0f0829ef | 270 | countries.append(object) |
6838b55a MT |
271 | |
272 | else: | |
92af07ad MT |
273 | log.warning("Invalid argument: %s" % object) |
274 | continue | |
275 | ||
276 | if not countries and not asns: | |
277 | log.error("Nothing to export") | |
278 | return 2 | |
6838b55a MT |
279 | |
280 | # Open the database | |
281 | try: | |
282 | db = location.Database(ns.database) | |
283 | except FileNotFoundError as e: | |
284 | log.error("Count not open database: %s" % ns.database) | |
285 | return 1 | |
286 | ||
287 | # Select the output format | |
288 | writer = self.output_formats.get(ns.format) | |
289 | assert writer | |
290 | ||
291 | e = Exporter(db, writer) | |
48bcc19d | 292 | e.export(ns.directory, countries=countries, asns=asns, families=families) |
6838b55a MT |
293 | |
294 | ||
295 | def main(): | |
296 | # Run the command line interface | |
297 | c = CLI() | |
298 | c.run() | |
299 | ||
300 | main() |