98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
"""dsc_datatool.generator.client_subnet_country
|
|
|
|
See `man dsc-datatool-generator client_subnet_country`.
|
|
|
|
Part of dsc_datatool.
|
|
|
|
:copyright: 2024 OARC, Inc.
|
|
"""
|
|
|
|
import maxminddb
|
|
import os
|
|
import logging
|
|
|
|
from dsc_datatool import Generator, Dataset, Dimension, args
|
|
|
|
|
|
class client_subnet_country(Generator):
|
|
reader = None
|
|
nonstrict = False
|
|
|
|
|
|
def __init__(self, opts):
|
|
Generator.__init__(self, opts)
|
|
paths = opts.get('path', ['/var/lib/GeoIP', '/usr/share/GeoIP', '/usr/local/share/GeoIP'])
|
|
if not isinstance(paths, list):
|
|
paths = [ paths ]
|
|
filename = opts.get('filename', 'GeoLite2-Country.mmdb')
|
|
db = opts.get('db', None)
|
|
|
|
if db is None:
|
|
for path in paths:
|
|
db = '%s/%s' % (path, filename)
|
|
if os.path.isfile(db) and os.access(db, os.R_OK):
|
|
break
|
|
db = None
|
|
if db is None:
|
|
raise Exception('Please specify valid Maxmind database with path=,filename= or db=')
|
|
|
|
logging.info('Using %s' % db)
|
|
self.reader = maxminddb.open_database(db)
|
|
|
|
if opts.get('nonstrict', False):
|
|
self.nonstrict = True
|
|
|
|
|
|
def process(self, datasets):
|
|
gen_datasets = []
|
|
|
|
for dataset in datasets:
|
|
if dataset.name != 'client_subnet':
|
|
continue
|
|
|
|
subnets = {}
|
|
for d1 in dataset.dimensions:
|
|
for d2 in d1.dimensions:
|
|
for k, v in d2.values.items():
|
|
if k == args.skipped_key:
|
|
continue
|
|
elif k == args.skipped_sum_key:
|
|
continue
|
|
|
|
if k in subnets:
|
|
subnets[k] += v
|
|
else:
|
|
subnets[k] = v
|
|
|
|
cc = {}
|
|
for subnet in subnets:
|
|
try:
|
|
c = self.reader.get(subnet)
|
|
except Exception as e:
|
|
if not self.nonstrict:
|
|
raise e
|
|
continue
|
|
if c:
|
|
iso_code = c.get('country', {}).get('iso_code', '??')
|
|
if iso_code in cc:
|
|
cc[iso_code] += subnets[subnet]
|
|
else:
|
|
cc[iso_code] = subnets[subnet]
|
|
|
|
if cc:
|
|
ccd = Dataset()
|
|
ccd.name = 'client_subnet_country'
|
|
ccd.start_time = dataset.start_time
|
|
ccd.stop_time = dataset.stop_time
|
|
gen_datasets.append(ccd)
|
|
|
|
ccd1 = Dimension('ClientCountry')
|
|
ccd1.values = cc
|
|
ccd.dimensions.append(ccd1)
|
|
|
|
return gen_datasets
|
|
|
|
|
|
import sys
|
|
if sys.version_info[0] == 3 and sys.version_info[1] == 5: # pragma: no cover
|
|
Generator.__init_subclass__(client_subnet_country)
|