Skip to content

Commit f220c04

Browse files
GalaxPGalaxP
committed
Added new module - geolite
1 parent 58eb6c5 commit f220c04

3 files changed

Lines changed: 256 additions & 0 deletions

File tree

‎geolite/Makefile.am‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
EXTRA_DIST=geolite.py readme.md
2+
bin_SCRIPTS=geolite.py
3+
4+
pkgdocdir=${docdir}/geolite
5+
pkgdoc_DATA=readme.md
6+
7+
pylint:
8+
pylint-3 geolite.py
9+
10+
flake8:
11+
flake8 geolite.py
12+
13+
pycodestyle:
14+
pycodestyle-3 geolite.py
15+
16+
lint: pylint flake8 pycodestyle
17+
18+
include ../aminclude.am

‎geolite/README.md‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# Geolite
2+
3+
## Module description
4+
5+
This module outputs flow records with geolocation data using a [geolite database](https://dev.maxmind.com/geoip/geolite2-free-geolocation-data/).
6+
7+
8+
## Input data
9+
10+
This module expects flow records in Unirec format. The required fields
11+
are determined by run time parameters.
12+
13+
14+
## Output data
15+
16+
Flows are sent on the output interface, also in Unirec format, they
17+
contain geolocation data. The fields included in the output interface will vary depending on the selected database type.
18+
19+
Below are the fields that will be set for each database type:
20+
* `country`
21+
* ipaddr `ip`
22+
* string `name`
23+
* string `iso_code`
24+
* uint32 `geoname_id`
25+
* uint32 `is_in_european_union`
26+
27+
* `city`
28+
* ipaddr `ip`
29+
* string `country_name`
30+
* string `country_iso_code`
31+
* uint32 `country_geoname_id`
32+
* uint32 `is_in_european_union`
33+
* string `city_name`
34+
* float `latitude`
35+
* float `longitude`
36+
* uint32 `accuracy_radius`
37+
38+
* `asn`
39+
* ipaddr `ip`
40+
* uint32 `asn`
41+
* string `string autonomous_system_organization`
42+
43+
44+
## Module parameters
45+
46+
In addition to the implicit *libtrap* parameters `-i IFC_SPEC`, `-h`
47+
and `-v` (see [Execute a
48+
module](https://github.com/CESNET/Nemea#try-out-nemea-modules)) this
49+
module takes the following parameters:
50+
51+
* `-d` `--db` path
52+
53+
* Specify path to the database file.
54+
55+
* `-f` `--fields` field1,field2,...
56+
57+
* Specify the name of field(s) from the input interface, which will be used for geolocation and lookup in the database (case sensitive).
58+
If multiple fields are specified, they must be separated by a comma.
59+
60+
* `-t` `--type` {country, city, asn}
61+
62+
* Specify the type of GeoLite database. The default value is `country`.
63+
64+
65+
## Example
66+
The following command :
67+
68+
`./geolite.py -i f:/etc/nemea/data/data.dan.trapcap,f:test.trapcap -d '/usr/share/GeoIP/GeoLite2-Country.mmdb' -t country -f "SRC_IP,DST_IP"`
69+
70+
will be interpreted as follows:
71+
72+
* `-i f:/etc/nemea/data/data.dan.trapcap,f:test.trapcap`
73+
sets the input and output interfaces to a file.
74+
75+
* `-d '/usr/share/GeoIP/GeoLite2-Country.mmdb'` sets the path to the database file.
76+
77+
* `-t country` sets the database type to `country` (can be omitted as it is the default).
78+
79+
* `-f "SRC_IP,DST_IP"` specifies the names of the fields containing IP addresses to be used for geolocation.
80+
81+
<!--- Local variables: -->
82+
<!--- mode: markdown; -->
83+
<!--- mode: auto-fill; -->
84+
<!--- mode: flyspell; -->
85+
<!--- ispell-local-dictionary: "british"; -->
86+
<!--- End: -->

‎geolite/geolite.py‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
#!/usr/bin/python3
2+
#
3+
# Copyright (C) 2025 CESNET
4+
#
5+
# LICENSE TERMS
6+
#
7+
# Redistribution and use in source and binary forms, with or without
8+
# modification, are permitted provided that the following conditions
9+
# are met:
10+
# 1. Redistributions of source code must retain the above copyright
11+
# notice, this list of conditions and the following disclaimer.
12+
# 2. Redistributions in binary form must reproduce the above copyright
13+
# notice, this list of conditions and the following disclaimer in
14+
# the documentation and/or other materials provided with the
15+
# distribution.
16+
# 3. Neither the name of the Company nor the names of its contributors
17+
# may be used to endorse or promote products derived from this
18+
# software without specific prior written permission.
19+
#
20+
# ALTERNATIVELY, provided that this notice is retained in full, this
21+
# product may be distributed under the terms of the GNU General Public
22+
# License (GPL) version 2 or later, in which case the provisions
23+
# of the GPL apply INSTEAD OF those given above.
24+
#
25+
# This software is provided ``as is'', and any express or implied
26+
# warranties, including, but not limited to, the implied warranties of
27+
# merchantability and fitness for a particular purpose are disclaimed.
28+
# In no event shall the company or contributors be liable for any
29+
# direct, indirect, incidental, special, exemplary, or consequential
30+
# damages (including, but not limited to, procurement of substitute
31+
# goods or services; loss of use, data, or profits; or business
32+
# interruption) however caused and on any theory of liability, whether
33+
# in contract, strict liability, or tort (including negligence or
34+
# otherwise) arising in any way out of the use of this software, even
35+
# if advised of the possibility of such damage.
36+
37+
38+
import argparse
39+
import pytrap
40+
import geoip2.database
41+
42+
43+
CITY_OUTPUTSPEC = "ipaddr ip, string country_name, string country_iso_code, uint32 country_geoname_id, uint32 is_in_european_union, string city_name, float latitude, float longitude, uint32 accuracy_radius"
44+
COUNTRY_OUTPUTSPEC = "ipaddr ip, string name, string iso_code, uint32 geoname_id, uint32 is_in_european_union"
45+
ASN_OUTPUTSPEC = "ipaddr ip, uint32 asn, string autonomous_system_organization"
46+
47+
48+
parser = argparse.ArgumentParser(description='Module for geolocation using GeoLite2 database')
49+
parser.add_argument('-i', "--ifcspec",
50+
help="select TRAP IFC specifier")
51+
parser.add_argument('-d', "--db",
52+
help="path to the GeoLite2 database file",
53+
required=True)
54+
parser.add_argument('-f', "--fields",
55+
help="input fields to use for geolocation seperated by comma",
56+
default="SRC_IP")
57+
parser.add_argument('-t', "--type",
58+
help="type of GeoLite database",
59+
choices=['country', 'city', 'asn'],
60+
default="country")
61+
62+
# parse command line arguments
63+
args = parser.parse_args()
64+
fields = args.fields.split(',')
65+
66+
67+
# initialize TRAP context
68+
trap = pytrap.TrapCtx()
69+
trap.init(['-i', args.ifcspec], 1, 1)
70+
71+
# output interface
72+
fmttype = pytrap.FMT_UNIREC
73+
74+
# set the correct output specification based on the type of GeoLite database
75+
if args.type == 'asn':
76+
outputspec = ASN_OUTPUTSPEC
77+
elif args.type == 'city':
78+
outputspec = CITY_OUTPUTSPEC
79+
else:
80+
outputspec = COUNTRY_OUTPUTSPEC
81+
82+
trap.setDataFmt(0, fmttype, outputspec)
83+
output = pytrap.UnirecTemplate(outputspec)
84+
output.createMessage()
85+
86+
# input interface
87+
fmtspec = ""
88+
trap.setRequiredFmt(0, fmttype, fmtspec)
89+
rec = pytrap.UnirecTemplate(fmtspec)
90+
91+
92+
# open GeoLite2 database reader
93+
with geoip2.database.Reader(args.db) as reader:
94+
# main loop
95+
while True:
96+
try:
97+
data = trap.recv()
98+
except pytrap.FormatChanged as e:
99+
fmttype, fmtspec = trap.getDataFmt(0)
100+
rec = pytrap.UnirecTemplate(fmtspec)
101+
data = e.data
102+
except pytrap.Terminated:
103+
print("Terminated trap.")
104+
break
105+
except pytrap.TrapError:
106+
print("Trap error, exiting.")
107+
break
108+
if len(data) <= 1:
109+
break
110+
111+
else:
112+
rec.setData(data)
113+
for field in fields:
114+
ip = rec.get(data, field)
115+
output.ip = ip
116+
try:
117+
if args.type == 'city':
118+
geolocation = reader.city(str(ip))
119+
elif args.type == 'country':
120+
geolocation = reader.country(str(ip))
121+
elif args.type == 'asn':
122+
geolocation = reader.asn(str(ip))
123+
except:
124+
continue
125+
if geolocation is None:
126+
continue
127+
128+
# fill output fields based on geolocation type
129+
if args.type == 'country':
130+
output.name = geolocation.country.name or "unkown"
131+
output.iso_code = geolocation.country.iso_code or "unkown"
132+
output.geoname_id = geolocation.country.geoname_id or 0
133+
output.is_in_european_union = geolocation.country.is_in_european_union or False
134+
elif args.type == 'city':
135+
output.country_name = geolocation.country.name or "unkown"
136+
output.country_iso_code = geolocation.country.iso_code or "unkown"
137+
output.country_geoname_id = geolocation.country.geoname_id or 0
138+
output.is_in_european_union = geolocation.country.is_in_european_union or False
139+
output.city_name = geolocation.city.name or "unknown"
140+
output.latitude = geolocation.location.latitude or 0.0
141+
output.longitude = geolocation.location.longitude or 0.0
142+
output.accuracy_radius = geolocation.location.accuracy_radius or 0
143+
elif args.type == 'asn':
144+
output.asn = geolocation.autonomous_system_number or 0
145+
output.autonomous_system_organization = geolocation.autonomous_system_organization or "unknown"
146+
147+
# send output data
148+
trap.send(output.getData(), 0)
149+
150+
# send end-of-stream message and exit
151+
trap.sendFlush(0)
152+
trap.finalize()

0 commit comments

Comments
 (0)