Merge aa0e1be into 0fefce4

dev-requirements.txt

@@ -1,6 +1,6 @@
-setuptools
-wheel
-twine
-coverage
-flake8
-sphinx
+setuptools~=47.3.1
+twine~=3.1.1
+wheel~=0.34.2
+coverage~=5.1
+flake8~=3.8.3
+sphinx~=3.1.1

devscripts/test.sh

@@ -4,6 +4,6 @@ flake8 libchickadee --count --show-source --statistics
 coverage run -m unittest discover
 coverage xml
 coverage report
-cd doc_src
+cd doc_src || exit
 make html
-cd ../devscripts
+cd ../devscripts || exit

libchickadee/backends/__init__.py

@@ -4,7 +4,7 @@
 
 A base class for handling resolution of IP addresses.
 
-This class includes elemends common across backend resolution data sources.
+This class includes elements common across backend resolution data sources.
 This includes common parameters, such as field names, and functions to handle
 querying the data sources. Additionally, this includes writers for CSV and JSON
 formats.
@@ -62,14 +62,15 @@ def query(self, data):
         Args:
             data (list, tuple, set, str): One or more IPs to resolve
 
-        Return:
-            (yield) request data iterator
+        Yield:
+            (dict) request data iterator
 
         Example:
             >>> records = ['1.1.1.1', '2.2.2.2']
             >>> resolver = ResolverBase()
             >>> resolved_data = resolver.query(records)
         """
+
         self.data = data
         if isinstance(data, (list, tuple, set)):
             return self.batch()
@@ -124,6 +125,7 @@ def write_json(outfile, data, headers=None, lines=False):
         Args:
             outfile (str or file_obj): Path to or already open file
             data (list): List of dictionaries containing resolved data
+            headers (list): List of column headers. Will use the first element of data if not present.
             lines (bool): Whether to export 1 dictionary object per line or
                 a whole json object.
 
@@ -145,26 +147,13 @@ def write_json(outfile, data, headers=None, lines=False):
 
         """
         was_opened = False
+        open_file = outfile
         if isinstance(outfile, str):
             open_file = open(outfile, 'w', newline="")
             was_opened = True
-        else:
-            open_file = outfile
 
         if headers:
-            # Only include fields in headers
-            # Include headers with no value if not present in original
-            selected_data = []
-            for x in data:
-                d = {}
-                for k, v in x.items():
-                    if k in headers:
-                        d[k] = v
-                for h in headers:
-                    if h not in d:
-                        d[h] = None
-                selected_data.append(d)
-            data = selected_data
+            data = ResolverBase.normalize_data_headers(data, headers)
 
         if lines:
             for entry in data:
@@ -174,3 +163,19 @@ def write_json(outfile, data, headers=None, lines=False):
 
         if was_opened:
             open_file.close()
+
+    @staticmethod
+    def normalize_data_headers(data, headers):
+        # Only include fields in headers
+        # Include headers with no value if not present in original
+        selected_data = []
+        for x in data:
+            d = {}
+            for k, v in x.items():
+                if k in headers:
+                    d[k] = v
+            for h in headers:
+                if h not in d:
+                    d[h] = None
+            selected_data.append(d)
+        return selected_data

libchickadee/backends/ipapi.py

@@ -134,7 +134,7 @@ def rate_limit(self, headers):
         duration of ``X-Ttl`` + 1 second.
 
         Args:
-            headers (dict): Request header information.
+            headers (dict, CaseInsensitiveDict): Request header information.
 
         Return:
             None
@@ -201,18 +201,14 @@ def batch(self):
                 self.rate_limit(rdata.headers)
                 result_list = [x for x in rdata.json()]
                 resolved_recs += result_list
-            elif rdata.status_code == 429:  # pragma: no cover
+            elif rdata.status_code == 429:
                 self.rate_limit(rdata.headers)
                 self.sleeper()
                 return self.batch()
             else:  # pragma: no cover
                 msg = "Unknown error encountered: {}".format(rdata.status_code)
                 logger.error(msg)
-                result_list = []
-                for result in records[x:x+100]:
-                    result_list.append({'query': result,
-                                        'status': 'failed',
-                                        'message': msg})
+                result_list = [{'query': result, 'status': 'failed', 'message': msg} for result in records[x:x+100]]
                 resolved_recs += result_list
         return resolved_recs
 
@@ -242,7 +238,7 @@ def single(self):
         if rdata.status_code == 200:
             self.rate_limit(rdata.headers)
             return rdata.json()
-        elif rdata.status_code == 429:  # pragma: no cover
+        elif rdata.status_code == 429:
             self.rate_limit(rdata.headers)
             self.sleeper()
             return self.single()
@@ -266,7 +262,7 @@ class ProResolver(Resolver):
     def __init__(self, api_key, fields=None, lang='en'):  # pragma: no cover
         if not fields:
             fields = FIELDS
-        Resolver.__init__(self, fields=fields, lang='en')
+        Resolver.__init__(self, fields=fields, lang=lang)
         self.uri = 'https://pro.ip-api.com/'
         self.api_key = api_key
         self.enable_sleep = False

libchickadee/chickadee.py

@@ -137,7 +137,6 @@
 import _io
 import configparser
 
-# Third Party Libs
 from tqdm import tqdm
 
 # Import lib features
@@ -221,7 +220,6 @@ def run(self, input_data, api_key=None):
             (list): List of dictionaries containing resolved hits.
         """
         self.input_data = input_data
-        results = []
         result_dict = {}
         # Extract and resolve IP addresses
         if not isinstance(self.input_data, _io.TextIOWrapper) and \
@@ -244,8 +242,7 @@ def run(self, input_data, api_key=None):
 
         # Resolve if requested
         if self.resolve_ips:
-            results = self.resolve(result_dict, api_key)
-            return results
+            return self.resolve(result_dict, api_key)
 
         return [{'query': k, 'count': v, 'message': 'No resolve'}
                 for k, v in result_dict.items()]
@@ -269,7 +266,7 @@ def str_handler(data):
         distinct IPs with their associated frequency count.
 
         Args:
-            data (str): raw input data from user
+            data (list, str): raw input data from user
 
         Return:
             data_dict (dict): dictionary of distinct IP addresses to resolve.
@@ -352,7 +349,7 @@ def dir_handler(self, folder_path):
 
     def resolve(self, data_dict, api_key=None):
         """Resolve IP addresses stored as keys within `data_dict`. The values
-        for each key should represent the number of occurances of an IP within
+        for each key should represent the number of occurrences of an IP within
         a data set.
 
         Args:
@@ -394,6 +391,7 @@ def resolve(self, data_dict, api_key=None):
             updated_results = []
             for result in results:
                 query = str(result.get('query', ''))
+                # noinspection PyTypeChecker
                 result['count'] = int(data_dict.get(query, '0'))
                 updated_results.append(result)
 
@@ -444,7 +442,7 @@ def setup_logging(path, verbose=False):  # pragma: no cover
     # Set default logger level to DEBUG. You can change this later
     logger.setLevel(logging.DEBUG)
 
-    # Logging formatter. Best to keep consistent for most usecases
+    # Logging formatter. Best to keep consistent for most use cases
     log_format = logging.Formatter(
         '%(asctime)s %(filename)s %(levelname)s %(module)s '
         '%(funcName)s:%(lineno)d - %(message)s')
@@ -506,74 +504,82 @@ def config_handing(config_file=None, search_conf_path=None):
     }
 
     if not config_file:
-        if not search_conf_path:
-            # Config file search path order:
-            # 1. Current directory
-            # 2. User home directory
-            # 3. System wide directory
-            # Needs to be named chickadee.ini or .chickadee.ini for detection.
-            search_conf_path = [os.path.abspath('.'), os.path.expanduser('~')]
-            if 'win32' in sys.platform:
-                search_conf_path.append(
-                    os.path.join(os.getenv('APPDATA'), 'chickadee'))
-                search_conf_path.append('C:\\ProgramData\\chickadee')
-            elif 'linux' in sys.platform or 'darwin' in sys.platform:
-                search_conf_path.append(
-                    os.path.expanduser('~/.config/chickadee'))
-                search_conf_path.append('/etc/chickadee')
-
-        for location in search_conf_path:
-            if not os.path.exists(location) or not os.path.isdir(location):
-                logger.debug(
-                    "Unable to access config file location {}.".format(
-                        location))
-            elif 'chickadee.ini' in os.listdir(location):
-                config_file = os.path.join(location, 'chickadee.ini')
-            elif '.chickadee.ini' in os.listdir(location):
-                config_file = os.path.join(location, '.chickadee.ini')
+        config_file = find_config_file(config_file, search_conf_path)
 
     fail_warn = 'Relying on argument defaults'
     if not config_file:
         logger.debug('Config file not found. {}'.format(fail_warn))
         return
 
-    if not os.path.exists(config_file) or not os.path.isfile(config_file):
+    if not (os.path.exists(config_file) and os.path.isfile(config_file)):
         logger.debug('Error accessing config file {}. {}'.format(
             config_file, fail_warn))
         return
 
     conf = configparser.ConfigParser()
     conf.read(config_file)
 
-    config = {}
+    return parse_config_sections(conf, section_defs)
 
-    for section in section_defs:
-        if section in conf:
-            for k, v in section_defs[section].items():
-                conf_section = conf[section]
-                conf_value = None
-                if isinstance(v, str):
-                    conf_value = conf_section.get(k)
-                elif isinstance(v, list):
-                    conf_value = conf_section.get(k).split(',')
-                elif isinstance(v, bool):
-                    conf_value = conf_section.getboolean(k)
-                elif isinstance(v, dict):
-                    conf_value = conf_section.get(k)
-                    # Set backend args through nested option
-                    for sk, sv, in v.get(conf_value, {}).items():
-                        config[sk] = conf_section[sv]
-                config[k] = conf_value
 
+def parse_config_sections(conf, section_defs):
+    config = {}
+    for section, value in section_defs.items():
+        if section not in conf:
+            continue
+        for k, v in value.items():
+            conf_section = conf[section]
+            conf_value = None
+            if isinstance(v, str):
+                conf_value = conf_section.get(k)
+            elif isinstance(v, list):
+                conf_value = conf_section.get(k).split(',')
+            elif isinstance(v, bool):
+                conf_value = conf_section.getboolean(k)
+            elif isinstance(v, dict):
+                conf_value = conf_section.get(k)
+                # Set backend args through nested option
+                for sk, sv, in v.get(conf_value, {}).items():
+                    config[sk] = conf_section[sv]
+            config[k] = conf_value
     return config
 
 
+def find_config_file(config_file, search_conf_path):
+    if not search_conf_path:
+        # Config file search path order:
+        # 1. Current directory
+        # 2. User home directory
+        # 3. System wide directory
+        # Needs to be named chickadee.ini or .chickadee.ini for detection.
+        search_conf_path = [os.path.abspath('.'), os.path.expanduser('~')]
+        if 'win32' in sys.platform:
+            search_conf_path.append(
+                os.path.join(os.getenv('APPDATA'), 'chickadee'))
+            search_conf_path.append('C:\\ProgramData\\chickadee')
+        elif 'linux' in sys.platform or 'darwin' in sys.platform:
+            search_conf_path.append(
+                os.path.expanduser('~/.config/chickadee'))
+            search_conf_path.append('/etc/chickadee')
+    for location in search_conf_path:
+        if not (os.path.exists(location) and os.path.isdir(location)):
+            logger.debug(
+                "Unable to access config file location {}.".format(
+                    location))
+        elif 'chickadee.ini' in os.listdir(location):
+            config_file = os.path.join(location, 'chickadee.ini')
+        elif '.chickadee.ini' in os.listdir(location):
+            config_file = os.path.join(location, '.chickadee.ini')
+    return config_file
+
+
 def arg_handling(args):
     """Parses command line arguments.
 
     Returns:
         argparse Namespace containing argument parameters.
     """
+    # noinspection PyTypeChecker
     parser = argparse.ArgumentParser(
         description=__desc__,
         formatter_class=CustomArgFormatter,
@@ -694,12 +700,14 @@ def join_config_args(config, args, definitions=None):
     return final_config
 
 
-def entry(args=sys.argv[1:]):  # pragma: no cover
+def entry(args=None):  # pragma: no cover
     """Entrypoint for package script.
 
     Args:
         args: Arguments from invocation.
     """
+    if not args:
+        args = sys.argv[1:]
     # Handle parameters from config file and command line.
     args = arg_handling(args)
     config = config_handing(args.config)

libchickadee/parsers/__init__.py

@@ -8,7 +8,7 @@
 for resolution.
 
 """
-
+import os
 import re
 
 from netaddr import IPAddress
@@ -60,10 +60,21 @@
 IPv6Pattern = re.compile(IPV6ADDR)
 
 
+def run_parser_from_cli(args, parser_obj):
+    if os.path.isdir(args.path):
+        for root, _, files in os.walk(args.path):
+            for fentry in files:
+                parser_obj.parse_file(os.path.join(root, fentry))
+    else:
+        parser_obj.parse_file(args.path)
+    print("{} unique IPs discovered".format(len(parser_obj.ips)))
+
+
 class ParserBase(object):
     """Base class for parsers, containing common utilities."""
     def __init__(self, ignore_bogon=True):
         self.ignore_bogon = ignore_bogon
+        self.ips = {}
 
     def check_ips(self, data):
         """Check data for IP addresses. Results stored in ``self.ips``.
@@ -115,7 +126,5 @@ def is_bogon(ip_addr):
             (bool): Whether or not the IP is a known BOGON address.
         """
         ip = IPAddress(ip_addr)
-        if (ip.is_private() or ip.is_link_local() or
-                ip.is_reserved() or ip.is_multicast()):
-            return True
-        return False
+        return bool((ip.is_private() or ip.is_link_local() or
+                     ip.is_reserved() or ip.is_multicast()))