Skip to content

Commit

Permalink
Fixed #27 and #28
Browse files Browse the repository at this point in the history 
Corrected bug in variable assignment and data type error.
Added unittest for no-resolve logic
  • Loading branch information
@chapinb
chapinb committed Jan 14, 2020
1 parent 2165ebe commit 6435cd0
Show file tree
Hide file tree
Showing 6 changed files with 42 additions and 35 deletions.
8 changes: 3 additions & 5 deletions libchickadee/backends/ipapi.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
logger = logging.getLogger('libchickadee.chickadee')

__author__ = 'Chapin Bryce'
__date__ = 20200107
__date__ = 20200114
__license__ = 'GPLv3 Copyright 2019 Chapin Bryce'
__desc__ = '''Yet another GeoIP resolution tool.'''

Expand Down Expand Up @@ -92,9 +92,7 @@ def batch(self):

if rdata.status_code == 200:
self.rate_limit(rdata.headers)
result_list = []
for result in rdata.json():
result_list.append(result)
result_list = [x for x in rdata.json()]
resolved_recs += result_list
elif rdata.status_code == 429:
self.rate_limit(rdata.headers)
Expand Down Expand Up @@ -129,7 +127,7 @@ def single(self):
)
if rdata.status_code == 200:
self.rate_limit(rdata.headers)
return [rdata.json()]
return rdata.json()
elif rdata.status_code == 429:
self.rate_limit(rdata.headers)
self.sleeper()
Expand Down
15 changes: 9 additions & 6 deletions libchickadee/chickadee.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@


__author__ = 'Chapin Bryce'
__date__ = 20200107
__date__ = 20200114
__license__ = 'GPLv3 Copyright 2019 Chapin Bryce'
__desc__ = '''Yet another GeoIP resolution tool.
Expand Down Expand Up @@ -54,7 +54,7 @@ def __init__(self, outformat='json', outfile=sys.stdout, fields=_FIELDS):
self.force_single = False
self.lang = 'en'
self.pbar = False
self.resolve_ips=True
self.resolve_ips = True

def run(self, input_data):
"""Evaluate the input data format to extract and resolve IP addresses.
Expand Down Expand Up @@ -110,8 +110,7 @@ def get_api_key():
api_key = os.environ.get('CHICKADEE_API_KEY', None)
if api_key is not None and len(api_key):
return api_key
else:
return None
return None

@staticmethod
def str_handler(data):
Expand Down Expand Up @@ -169,6 +168,7 @@ def resolve(self, data_dict):
logger.info("Resolving IPs")
if self.force_single:
results = []
data = distinct_ips
if self.pbar:
data = tqdm(distinct_ips, desc="Resolving IPs", unit_scale=True)

Expand All @@ -184,8 +184,11 @@ def resolve(self, data_dict):
if 'count' in self.fields:
updated_results = []
for result in results:
query = str(result.get('query', ''))
result['count'] = int(data_dict.get(query, 0))
try:
query = str(result.get('query', ''))
except AttributeError:
import pdb; pdb.set_trace()
result['count'] = int(data_dict.get(query, '0'))
updated_results.append(result)

return updated_results
Expand Down
4 changes: 2 additions & 2 deletions libchickadee/parsers/plain_text.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from libchickadee.parsers import IPv4Pattern, IPv6Pattern, strip_ipv6

__author__ = 'Chapin Bryce'
__date__ = 20200107
__date__ = 20200114
__license__ = 'GPLv3 Copyright 2019 Chapin Bryce'
__desc__ = '''Yet another GeoIP resolution tool.'''

Expand All @@ -36,7 +36,7 @@ def parse_file(self, file_entry, is_stream=False):
else:
file_data = open(file_entry, 'rb')
else:
if binascii.hexlify(file_entry.buffer.read(2)) == b'1f8b':
if binascii.hexlify(file_entry.buffer.read(2)) == b'1f8b':
file_data = GzipFile(fileobj=file_entry)
else:
file_data = file_entry.buffer
Expand Down
15 changes: 5 additions & 10 deletions libchickadee/test/test_backend_ipapi.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from libchickadee.backends.ipapi import Resolver

__author__ = 'Chapin Bryce'
__date__ = 20200107
__date__ = 20200114
__license__ = 'GPLv3 Copyright 2019 Chapin Bryce'
__desc__ = '''Yet another GeoIP resolution tool.'''

Expand All @@ -31,8 +31,7 @@ def test_ipapi_resolve_query_single(self):
"""Query Method Test"""
for count, ip in enumerate(self.test_data_ips):
data = self.resolver.query(ip)
res = [x for x in data]
self.assertEqual(res, [self.expected_result[count]])
self.assertEqual(data, self.expected_result[count])

def test_ipapi_resolve_query_batch(self):
"""Batch Query Method Test"""
Expand All @@ -50,8 +49,7 @@ def test_ipapi_resolve_single(self):
for count, ip in enumerate(self.test_data_ips):
self.resolver.data = ip
data = self.resolver.single()
res = [x for x in data]
self.assertEqual(res, [self.expected_result[count]])
self.assertEqual(data, self.expected_result[count])

def test_ipapi_resolve_batch(self):
"""Batch Query Method Test"""
Expand All @@ -60,22 +58,19 @@ def test_ipapi_resolve_batch(self):
res = [x for x in data]
self.assertCountEqual(res, self.expected_result)


def test_ipapi_resolve_single_field(self):
"""Single Query Method Test"""
for count, ip in enumerate(self.test_data_ips):
self.resolver.data = ip
self.resolver.fields = ['query', 'country', 'as']
data = self.resolver.single()
res = [x for x in data]

expected = {}
for field in ['query', 'country', 'as']:
if field not in res[0]:
if field not in data:
continue
expected[field] = self.expected_result[count].get(field, None)
self.assertEqual(res, [expected])

self.assertEqual(data, expected)

'''
import os
Expand Down
33 changes: 22 additions & 11 deletions libchickadee/test/test_chickadee.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from libchickadee.chickadee import Chickadee

__author__ = 'Chapin Bryce'
__date__ = 20200107
__date__ = 20200114
__license__ = 'GPLv3 Copyright 2019 Chapin Bryce'
__desc__ = '''Yet another GeoIP resolution tool.'''

Expand All @@ -28,7 +28,20 @@ def setUp(self):
]

self.fields = ['query', 'count', 'as', 'country', 'org', 'proxy']
self.maxDiff = None

def test_no_resolve(self):
results = [
{'query': '10.0.1.2', 'count': 1, 'message': 'No resolve'},
{'query': '8.8.8.8', 'count': 1, 'message': 'No resolve'},
{'query': '2001:4860:4860::8888', 'count': 1, 'message': 'No resolve'}
]
for count, ip in enumerate(self.test_data_ips):
chickadee = Chickadee()
chickadee.fields = self.fields
chickadee.resolve_ips = False
data = chickadee.run(ip)
res = [x for x in data]
self.assertEqual(res, [results[count]])

def test_chickadee_single(self):
"""Query Method Test"""
Expand Down Expand Up @@ -94,7 +107,6 @@ def setUp(self):
]

self.fields = ['query', 'count', 'as', 'country', 'org', 'proxy']
self.maxDiff = None

self.xlsx_data_results = [
{"as": "AS15169 Google LLC",
Expand Down Expand Up @@ -169,38 +181,37 @@ def test_ipapi_resolve_query_folder(self):
"""Batch Query Method Test"""
expected = [
{"country": "Australia", "org": "", "as": "AS13335 Cloudflare, Inc.",
"proxy": False, "query": "1.1.1.1", "count": 6},
"proxy": False, "query": "1.1.1.1", "count": 6},

{"query": "10.0.1.2", "count": 3},

{"country": "United States", "org": "Level 3", "as": "AS15169 Google LLC",
"proxy": False, "query": "8.8.8.8", "count": 3},
"proxy": False, "query": "8.8.8.8", "count": 3},

{"country": "United States", "org": "Google LLC", "as": "AS15169 Google LLC",
"proxy": False, "query": "2001:4860:4860::8888", "count": 3},
"proxy": False, "query": "2001:4860:4860::8888", "count": 3},

{"country": "United States", "org": "Informs",
"as": "AS3356 Level 3 Communications, Inc.", "proxy": True, "query": "4.4.4.4",
"count": 3},

{"country": "United States", "org": "Google LLC", "as": "AS15169 Google LLC",
"proxy": False, "query": "2001:4860:4860::8844", "count": 4},
"proxy": False, "query": "2001:4860:4860::8844", "count": 4},

{"country": "United States", "org": "Google LLC", "as": "AS15169 Google LLC",
"proxy": False, "query": "2001:4860:4860::8888", "count": 3},
"proxy": False, "query": "2001:4860:4860::8888", "count": 3},

{"country": "France", "org": "", "as": "AS3215 Orange S.A.",
"proxy": True, "query": "2.2.2.2", "count": 3},
"proxy": True, "query": "2.2.2.2", "count": 3},

{"country": "United States", "org": "Google LLC", "as": "AS15169 Google LLC",
"proxy": False, "query": "2001:4860:4860::8844", "count": 4}
"proxy": False, "query": "2001:4860:4860::8844", "count": 4}
]

chickadee = Chickadee()
chickadee.fields = self.fields
data = chickadee.run(self.test_data_dir)
res = [x for x in data]
batch_result = [] # No reverse field
self.assertCountEqual(res, expected)

if __name__ == '__main__':
Expand Down
2 changes: 1 addition & 1 deletion setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

setuptools.setup(
name='chickadee',
version=20200107,
version=20200114,
description='Yet another GeoIP resolution tool.',
author='Chapin Bryce',
author_email='python@chapinb.com',
Expand Down

0 comments on commit 6435cd0

Please sign in to comment.