def extract_ips(c, filename):
"""Extracts all IPv4 ip addresses out of @filename
"""
with open(filename, mode='r') as fd:
content = fd.read()
sorted_ips = sort_ips([t.value for t in create_lexer(content)])
tables = [
models.GeoLite2AsnBlocksIpv4,
models.GeoLite2CityBlocksIpv4,
]
with open('results.json', mode='w') as fd:
for ip in sorted_ips:
network = '.'.join(ip.split('.')[0:3])
network_wildcard = '{}%'.format(network)
with sqlite3.session_scope() as session:
for table in tables:
results = session.query(table).filter(
table.network.startswith(network_wildcard)) # noqa
if results.all():
json.dump([x.to_dict() for x in results.all()],
fp=fd,
ensure_ascii=False,
separators=(',', ': ')) # noqa
def populate_sqlite3(c):
"""Populate SQLite3 db with geolite2 CSV data
"""
with c.cd(PROJECT_ROOT_DIR):
def geolite2_city_blocks_ipv4():
return models.GeoLite2CityBlocksIpv4
def geolite2_asn_blocks_ipv4():
return models.GeoLite2AsnBlocksIpv4
# create a mapping of CSV files to models but to avoid the import
# problem we just put this behind a small wrapper function.
model_mapping = {
'GeoLite2-City-Blocks-IPv4.csv': geolite2_city_blocks_ipv4,
'GeoLite2-ASN-Blocks-IPv4.csv': geolite2_asn_blocks_ipv4,
}
csv_files = [
os.path.join(
'data',
'geolite2',
'city',
'GeoLite2-City-Blocks-IPv4.csv',
),
os.path.join(
'data',
'geolite2',
'asn',
'GeoLite2-ASN-Blocks-IPv4.csv',
)
]
for csv_file in csv_files:
reader = csv.reader(open(csv_file))
headers = next(reader)
ModelClass = model_mapping[os.path.basename(csv_file)]()
with sqlite3.session_scope() as session:
for i, line in enumerate(reader):
model = ModelClass()
for key, value in zip(headers, line):
setattr(model, key, value)
log.debug('adding model={}'.format(model.__dict__))
session.add(model)
if i % 10000 == 0:
# commit every 10000 records to avoid OOM
session.commit()
def on_call(db_filename=None, **kwargs):
kwargs.setdefault('echo', True)
db_filename = db_filename or 'test.sqlite3'
# we must recreate the Session object since we are using a global
# static variable. We do this so that one test that modifies the db
# doesn't cause another test to use the same session. We only care
# about this here since each test should be isolated.
sqlite3.Session = sqlite3.sessionmaker()
engine = sqlite3.init_engine(filename=db_filename)
sqlite3.init_db(engine=engine)
return sqlite3.session_scope()
def test_when_session_is_created_without_initializing_the_db_we_attempt_to_recover_and_initialize_the_db_and_recreate_a_new_session( # noqa
self, tmpdir):
tmpdir.chdir()
data = {
'network': '223.255.254.0/24',
'autonomous_system_number': 55415,
'autonomous_system_organization': 'MARINA BAY SANDS PTE LTD'
}
rec = models.GeoLite2AsnBlocksIpv4(**data)
# we must recreate the session
sqlite3.Session = sqlite3.sessionmaker()
with sqlite3.session_scope() as session:
session.add(rec)
results = session.query(models.GeoLite2AsnBlocksIpv4)
assert len(results.all()) == 1
def test_attempting_to_save_a_record_to_the_db_that_raises_an_exception_we_should_rollback_the_transaction( # noqa
self, ):
# we create a dummy object to represent a sessionmaker instance
# we will replace the real implementation with this one
class MockedSession(object):
bind = None
def add(*args, **kwargs):
pass
def commit(*args, **kwargs):
raise Exception('boom')
def rollback(*args, **kwargs):
pass
def close(*args, **kwargs):
pass
with mock.patch.object(
sqlite3,
'Session',
return_value=MockedSession,
):
data = {
'network': '223.255.254.0/24',
'autonomous_system_number': 55415,
'autonomous_system_organization': 'MARINA BAY SANDS PTE LTD'
}
rec = models.GeoLite2AsnBlocksIpv4(**data)
with pytest.raises(Exception) as exp:
with sqlite3.session_scope() as session:
session.add(rec)
assert 'boom' in str(exp.value)