def test_empty_file(self):
path = 'test.txt'
open(path, 'w').close()
expected = {}
actual = reader.read_ips(path)
self.assertEqual(expected, actual)
os.remove(path)
def test(self):
expected = {
'172.217.12.14': 1,
'157.240.18.35': 2,
'72.30.35.10': 1,
'151.101.64.81': 1,
'151.101.1.164': 1,
'208.80.153.224': 2,
'192.30.253.113': 1,
'108.174.10.10': 2,
'18.205.93.0': 1,
'151.101.193.69': 1
}
actual = reader.read_ips('tests/resources/ips.txt')
self.assertEqual(expected, actual)
def read_data(self, path):
"""
Read file, pull out IPs, retrieve GeoIP and RDAP info on them, and write data to disk.
:param path: path of file containing IPs
:return: None
"""
print('\nReading IPs from file...')
ips = reader.read_ips(path)
num_ips = len(ips)
print(f'{num_ips} IPs found.')
print('\nRetrieving GeoIP and RDAP data for IPs...')
print('0%')
with Manager() as manager:
original_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
pool = Pool(4)
signal.signal(signal.SIGINT, original_sigint_handler)
shared = manager.list([{'i': 0, 'p': 0, 'num': num_ips}])
retrieve = partial(self.retrieve_ip_data, shared)
try:
results = pool.map_async(retrieve, ips.keys())
results = results.get(3600)
except KeyboardInterrupt:
print('Terminating')
pool.terminate()
sys.exit(0)
else:
pool.close()
pool.join()
print('Done.')
print('Writing data to disk...')
write_stats = {
'geoip': {
'added': 0,
'skipped': 0
},
'rdap': {
'added': 0,
'skipped': 0
},
'ip_rdap': {
'added': 0,
'skipped': 0
}
}
with self.warehouse.open() as wh:
for result in results:
ip = result.pop('ip')
for index, data in result.items():
if not data:
continue
if index == 'geoip':
added = wh.write(index, data)
if added:
write_stats[index]['added'] += 1
else:
write_stats[index]['skipped'] += 1
else:
for item in data:
if index == 'ip_rdap':
added = wh.write(index, {'ip': ip, 'handle': item['handle']})
else:
added = wh.write(index, item)
if added:
write_stats[index]['added'] += 1
else:
write_stats[index]['skipped'] += 1
for index, stats in write_stats.items():
print(f'\n{index}')
added = stats['added']
skipped = stats['skipped']
print(f'added: {added}, skipped: {skipped}')
def test_none_path(self):
with self.assertRaises(TypeError):
reader.read_ips(None)
def test_wrong_path(self):
with self.assertRaises(FileNotFoundError):
reader.read_ips('path/does/not/exist.txt')