BGP-Ranking/bgpranking/dbinsert.py

105 lines
4.9 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import logging
from redis import StrictRedis
from .libs.helpers import shutdown_requested, set_running, unset_running, get_socket_path, get_ipasn, sanity_check_ipasn
class DatabaseInsert():
def __init__(self, loglevel: int=logging.DEBUG):
self.__init_logger(loglevel)
self.ardb_storage = StrictRedis(unix_socket_path=get_socket_path('storage'), decode_responses=True)
self.redis_sanitized = StrictRedis(unix_socket_path=get_socket_path('prepare'), db=0, decode_responses=True)
self.ipasn = get_ipasn()
self.logger.debug('Starting import')
def __init_logger(self, loglevel):
self.logger = logging.getLogger(f'{self.__class__.__name__}')
self.logger.setLevel(loglevel)
def insert(self):
ready, message = sanity_check_ipasn(self.ipasn)
if not ready:
# Try again later.
self.logger.warning(message)
return
self.logger.debug(message)
set_running(self.__class__.__name__)
while True:
if shutdown_requested() or not self.ipasn.is_up:
break
uuids = self.redis_sanitized.spop('to_insert', 100)
if not uuids:
break
p = self.redis_sanitized.pipeline(transaction=False)
[p.hgetall(uuid) for uuid in uuids]
sanitized_data = p.execute()
for_query = []
for i, uuid in enumerate(uuids):
data = sanitized_data[i]
if not data:
self.logger.warning(f'No data for UUID {uuid}. This should not happen, but lets move on.')
continue
for_query.append({'ip': data['ip'], 'address_family': data['address_family'], 'source': 'caida',
'date': data['datetime'], 'precision_delta': {'days': 3}})
try:
responses = self.ipasn.mass_query(for_query)
except Exception:
self.logger.exception('Mass query in IPASN History failed, trying again later.')
# Rollback the spop
self.redis_sanitized.sadd('to_insert', *uuids)
break
retry = []
done = []
ardb_pipeline = self.ardb_storage.pipeline(transaction=False)
for i, uuid in enumerate(uuids):
data = sanitized_data[i]
if not data:
self.logger.warning(f'No data for UUID {uuid}. This should not happen, but lets move on.')
continue
routing_info = responses['responses'][i]['response'] # our queries are on one single date, not a range
# Data gathered from IPASN History:
# * IP Block of the IP
# * AS number
if 'error' in routing_info:
self.logger.warning(f"Unable to find routing information for {data['ip']} - {data['datetime']}: {routing_info['error']}")
continue
# Single date query, getting from the object
datetime_routing = list(routing_info.keys())[0]
entry = routing_info[datetime_routing]
if not entry:
# routing info is missing, need to try again later.
retry.append(uuid)
continue
if 'asn' in entry and entry['asn'] is None:
self.logger.warning(f"Unable to find the AS number associated to {data['ip']} - {data['datetime']} (got None). This should not happen...")
continue
if 'prefix' in entry and entry['prefix'] is None:
self.logger.warning(f"Unable to find the prefix associated to {data['ip']} - {data['datetime']} (got None). This should not happen...")
continue
# Format: <YYYY-MM-DD>|sources -> set([<source>, ...])
ardb_pipeline.sadd(f"{data['date']}|sources", data['source'])
# Format: <YYYY-MM-DD>|<source> -> set([<asn>, ...])
ardb_pipeline.sadd(f"{data['date']}|{data['source']}", entry['asn'])
# Format: <YYYY-MM-DD>|<source>|<asn> -> set([<prefix>, ...])
ardb_pipeline.sadd(f"{data['date']}|{data['source']}|{entry['asn']}", entry['prefix'])
# Format: <YYYY-MM-DD>|<source>|<asn>|<prefix> -> set([<ip>|<datetime>, ...])
ardb_pipeline.sadd(f"{data['date']}|{data['source']}|{entry['asn']}|{entry['prefix']}",
f"{data['ip']}|{data['datetime']}")
done.append(uuid)
ardb_pipeline.execute()
p = self.redis_sanitized.pipeline(transaction=False)
if done:
p.delete(*done)
if retry:
p.sadd('to_insert', *retry)
p.execute()
unset_running(self.__class__.__name__)