diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 714a09e..c6df0ed 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -1,6 +1,11 @@ name: Test and Publish ipdata to PyPI -on: push +on: + push: + branches: + - master + tags: + - '*' jobs: build-n-publish: diff --git a/README.md b/README.md index f59dcfa..3918dc8 100644 --- a/README.md +++ b/README.md @@ -57,10 +57,10 @@ Install the latest version of the cli with `pip`. pip install ipdata ``` -or `easy_install` +To use the [IPTrie](#iptrie) data structure, install with the `trie` extra (requires a C compiler): ```bash -easy_install ipdata +pip install ipdata[trie] ``` ## Library Usage @@ -580,6 +580,8 @@ IPTrie is a production-ready, type-safe trie for IP addresses and CIDR prefixes ### Quick Start +> **Note:** IPTrie requires the `trie` extra: `pip install ipdata[trie]` + ```python from ipdata import IPTrie @@ -741,6 +743,12 @@ A list of possible errors is available at [Status Codes](https://docs.ipdata.co/ ## Tests +Install test dependencies: + +```shell +pip install ipdata[test,trie] +``` + To run all tests ```shell diff --git a/setup.cfg b/setup.cfg index df6fb94..8bcf0d2 100644 --- a/setup.cfg +++ b/setup.cfg @@ -29,9 +29,13 @@ install_requires = click click_default_group pyperclip + +[options.extras_require] +test = pytest hypothesis python-dotenv +trie = pytricia [options.entry_points] diff --git a/src/ipdata/__init__.py b/src/ipdata/__init__.py index 8d521de..d59215a 100644 --- a/src/ipdata/__init__.py +++ b/src/ipdata/__init__.py @@ -7,7 +7,11 @@ >>> ipdata.lookup() # or ipdata.lookup("8.8.8.8") """ from .ipdata import IPData -from .iptrie import IPTrie + +try: + from .iptrie import IPTrie +except ImportError: + IPTrie = None # Configuration api_key = None @@ -15,8 +19,8 @@ default_client = None -def lookup(resource="", fields=[]): - return _proxy("lookup", resource=resource, fields=fields) +def lookup(resource="", fields=[], select_field=None): + return _proxy("lookup", resource=resource, fields=fields, select_field=select_field) def bulk(resources, fields=[]): diff --git a/src/ipdata/cli.py b/src/ipdata/cli.py index 07ed0a0..d71b67e 100644 --- a/src/ipdata/cli.py +++ b/src/ipdata/cli.py @@ -46,7 +46,6 @@ ╰───────────────╯ """ import csv -import io import json import logging import os @@ -66,9 +65,8 @@ from rich.progress import Progress from rich.tree import Tree -from .lolcat import LolCat from .geofeeds import Geofeed, GeofeedValidationError -from .ipdata import DotDict, IPData +from .ipdata import IPData console = Console() @@ -76,7 +74,7 @@ logging.basicConfig( level="ERROR", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()] ) -log = logging.getLogger("rich") +log = logging.getLogger(__name__) API_KEY_FILE = f"{Path.home()}/.ipdata" @@ -95,19 +93,24 @@ def _lookup(ipdata, *args, **kwargs): def print_ascii_logo(): """ - Print cool ascii logo with lolcat. + Print cool ascii logo with rainbow colors. """ - options = DotDict({"animate": False, "os": 6, "spread": 3.0, "freq": 0.1}) - logo = """ - _ _ _ -(_)_ __ __| | __ _| |_ __ _ + from rich.text import Text + + logo = r""" + _ _ _ +(_)_ __ __| | __ _| |_ __ _ | | '_ \ / _` |/ _` | __/ _` | | | |_) | (_| | (_| | || (_| | |_| .__/ \__,_|\__,_|\__\__,_| |_| - """ - lol = LolCat() - lol.cat(io.StringIO(logo), options) +""" + rainbow = ["red", "orange1", "yellow", "green", "blue", "dark_violet"] + lines = logo.strip("\n").split("\n") + for i, line in enumerate(lines): + text = Text(line) + text.stylize(rainbow[i % len(rainbow)]) + console.print(text) def pretty_print_data(data): diff --git a/src/ipdata/geofeeds.py b/src/ipdata/geofeeds.py index f614e3f..7d8bcba 100644 --- a/src/ipdata/geofeeds.py +++ b/src/ipdata/geofeeds.py @@ -9,15 +9,11 @@ from pathlib import Path import requests -from rich.logging import RichHandler from .codes import COUNTRIES, REGION_CODES -FORMAT = "%(message)s" -logging.basicConfig( - level="ERROR", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()] -) -log = logging.getLogger("rich") +log = logging.getLogger(__name__) +log.addHandler(logging.NullHandler()) pwd = Path(__file__).parent.resolve() diff --git a/src/ipdata/ipdata.py b/src/ipdata/ipdata.py index 26b0817..b840d8c 100644 --- a/src/ipdata/ipdata.py +++ b/src/ipdata/ipdata.py @@ -26,12 +26,9 @@ import functools from requests.adapters import HTTPAdapter, Retry -from rich.logging import RichHandler -FORMAT = "%(message)s" -logging.basicConfig( - level="ERROR", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()] -) +logger = logging.getLogger(__name__) +logger.addHandler(logging.NullHandler()) class IPDataException(Exception): @@ -68,7 +65,7 @@ class IPData(object): :param debug: A boolean used to set the log level. Set to True when debugging. """ - log = logging.getLogger("rich") + log = logging.getLogger(__name__) valid_fields = { "ip", @@ -100,13 +97,15 @@ class IPData(object): def __init__( self, - api_key=os.environ.get("IPDATA_API_KEY"), + api_key=None, endpoint="https://api.ipdata.co/", timeout=60, retry_limit=7, retry_backoff_factor=1, debug=False, ): + if api_key is None: + api_key = os.environ.get("IPDATA_API_KEY") if not api_key: raise IPDataException("API Key not set. Set an API key via the 'IPDATA_API_KEY' environment variable or see the docs for other ways to do so.") # Request settings @@ -119,6 +118,13 @@ def __init__( # Enable debugging if debug: self.log.setLevel(logging.DEBUG) + if not any( + not isinstance(h, logging.NullHandler) for h in self.log.handlers + ): + from rich.logging import RichHandler + handler = RichHandler() + handler.setFormatter(logging.Formatter("%(message)s", datefmt="[%X]")) + self.log.addHandler(handler) # Work around renamed argument in urllib3. if hasattr(urllib3.util.Retry.DEFAULT, "allowed_methods"): @@ -167,17 +173,18 @@ def _validate_ip_address(self, ip): if request_ip.is_private or request_ip.is_reserved or request_ip.is_multicast: raise ValueError(f"{ip} is a reserved IP Address") - def lookup(self, resource="", fields=[]): + def lookup(self, resource="", fields=[], select_field=None): """ Makes a GET request to the IPData API for the specified 'resource' and the given 'fields'. :param resource: Either an IP address or an ASN prefixed by "AS" eg. "AS15169" :param fields: A collection of API fields to be returned + :param select_field: A single field name to return (convenience alternative to fields) :returns: An API response as a DotDict object to allow dot notation access of fields eg. data.ip, data.company.name, data.threat.blocklists[0].name etc :raises IPDataException: if the API call fails or if there is a failure in decoding the response. - :raises ValueError: if 'resource' is not a string + :raises ValueError: if 'resource' is not a string, or if both 'select_field' and 'fields' are provided """ if type(resource) is not str: raise ValueError(f"{resource} must be of type 'str'") @@ -189,6 +196,13 @@ def lookup(self, resource="", fields=[]): if resource and not resource.startswith("AS"): self._validate_ip_address(resource) + if select_field is not None: + if fields: + raise ValueError("'select_field' and 'fields' are mutually exclusive. Use one or the other.") + if not isinstance(select_field, str): + raise ValueError("'select_field' must be of type 'str'.") + fields = [select_field] + self._validate_fields(fields) query_params = self._query_params | {"fields": ",".join(fields)} diff --git a/src/ipdata/iptrie.py b/src/ipdata/iptrie.py index c6b6102..edbe906 100644 --- a/src/ipdata/iptrie.py +++ b/src/ipdata/iptrie.py @@ -20,7 +20,10 @@ from collections.abc import Iterator from typing import TypeVar, Generic -from pytricia import PyTricia +try: + from pytricia import PyTricia +except ImportError: + PyTricia = None T = TypeVar("T") @@ -94,6 +97,11 @@ class IPTrie(Generic[T]): def __init__(self) -> None: """Initialize an empty IPTrie.""" + if PyTricia is None: + raise ImportError( + "pytricia is required for IPTrie but is not installed. " + "Install it with: pip install ipdata[trie]" + ) self._ipv4: PyTricia = PyTricia(self.IPV4_BITS) self._ipv6: PyTricia = PyTricia(self.IPV6_BITS) diff --git a/src/ipdata/lolcat.py b/src/ipdata/lolcat.py deleted file mode 100644 index 424d034..0000000 --- a/src/ipdata/lolcat.py +++ /dev/null @@ -1,290 +0,0 @@ -#!/usr/bin/env python -# -# "THE BEER-WARE LICENSE" (Revision 43~maze) -# -# wrote these files. As long as you retain this notice you -# can do whatever you want with this stuff. If we meet some day, and you think -# this stuff is worth it, you can buy me a beer in return. - -from __future__ import print_function - -import atexit -import math -import os -import random -import re -import sys -import time -from signal import signal, SIGPIPE, SIG_DFL - -PY3 = sys.version_info >= (3,) - -# override default handler so no exceptions on SIGPIPE -signal(SIGPIPE, SIG_DFL) - -# Reset terminal colors at exit -def reset(): - sys.stdout.write("\x1b[0m") - sys.stdout.flush() - - -atexit.register(reset) - - -STRIP_ANSI = re.compile(r"\x1b\[(\d+)(;\d+)?(;\d+)?[m|K]") -COLOR_ANSI = ( - (0x00, 0x00, 0x00), - (0xCD, 0x00, 0x00), - (0x00, 0xCD, 0x00), - (0xCD, 0xCD, 0x00), - (0x00, 0x00, 0xEE), - (0xCD, 0x00, 0xCD), - (0x00, 0xCD, 0xCD), - (0xE5, 0xE5, 0xE5), - (0x7F, 0x7F, 0x7F), - (0xFF, 0x00, 0x00), - (0x00, 0xFF, 0x00), - (0xFF, 0xFF, 0x00), - (0x5C, 0x5C, 0xFF), - (0xFF, 0x00, 0xFF), - (0x00, 0xFF, 0xFF), - (0xFF, 0xFF, 0xFF), -) - - -class stdoutWin: - def __init__(self): - self.output = sys.stdout - self.string = "" - self.i = 0 - - def isatty(self): - return self.output.isatty() - - def write(self, s): - self.string = self.string + s - - def flush(self): - return self.output.flush() - - def prints(self): - string = 'echo|set /p="%s"' % (self.string) - os.system(string) - self.i += 1 - self.string = "" - - def println(self): - print() - self.prints() - - -class LolCat(object): - def __init__(self, mode=256, output=sys.stdout): - self.mode = mode - self.output = output - - def _distance(self, rgb1, rgb2): - return sum(map(lambda c: (c[0] - c[1]) ** 2, zip(rgb1, rgb2))) - - def ansi(self, rgb): - r, g, b = rgb - - if self.mode in (8, 16): - colors = COLOR_ANSI[: self.mode] - matches = [ - (self._distance(c, map(int, rgb)), i) for i, c in enumerate(colors) - ] - matches.sort() - color = matches[0][1] - - return "3%d" % (color,) - else: - gray_possible = True - sep = 2.5 - - while gray_possible: - if r < sep or g < sep or b < sep: - gray = r < sep and g < sep and b < sep - gray_possible = False - - sep += 42.5 - - if gray: - color = 232 + int(float(sum(rgb) / 33.0)) - else: - color = sum( - [16] - + [ - int(6 * float(val) / 256) * mod - for val, mod in zip(rgb, [36, 6, 1]) - ] - ) - - return "38;5;%d" % (color,) - - def wrap(self, *codes): - return "\x1b[%sm" % ("".join(codes),) - - def rainbow(self, freq, i): - r = math.sin(freq * i) * 127 + 128 - g = math.sin(freq * i + 2 * math.pi / 3) * 127 + 128 - b = math.sin(freq * i + 4 * math.pi / 3) * 127 + 128 - return [r, g, b] - - def cat(self, fd, options): - if options.animate: - self.output.write("\x1b[?25l") - - for line in fd: - options.os += 1 - self.println(line, options) - - if options.animate: - self.output.write("\x1b[?25h") - - def println(self, s, options): - s = s.rstrip() - if options.force or self.output.isatty(): - s = STRIP_ANSI.sub("", s) - - if options.animate: - self.println_ani(s, options) - else: - self.println_plain(s, options) - - self.output.write("\n") - self.output.flush() - if os.name == "nt": - self.output.println() - - def println_ani(self, s, options): - if not s: - return - - for i in range(1, options.duration): - self.output.write("\x1b[%dD" % (len(s),)) - self.output.flush() - options.os += options.spread - self.println_plain(s, options) - time.sleep(1.0 / options.speed) - - def println_plain(self, s, options): - for i, c in enumerate(s if PY3 else s.decode(options.charset_py2, "replace")): - rgb = self.rainbow(options.freq, options.os + i / options.spread) - self.output.write( - "".join( - [ - self.wrap(self.ansi(rgb)), - c if PY3 else c.encode(options.charset_py2, "replace"), - ] - ) - ) - if os.name == "nt": - self.output.print() - - -def detect_mode(term_hint="xterm-256color"): - """ - Poor-mans color mode detection. - """ - if "ANSICON" in os.environ: - return 16 - elif os.environ.get("ConEmuANSI", "OFF") == "ON": - return 256 - else: - term = os.environ.get("TERM", term_hint) - if term.endswith("-256color") or term in ("xterm", "screen"): - return 256 - elif term.endswith("-color") or term in ("rxvt",): - return 16 - else: - return 256 # optimistic default - - -def run(): - """Main entry point.""" - import optparse - - parser = optparse.OptionParser(usage=r"%prog [] [file ...]") - parser.add_option( - "-p", "--spread", type="float", default=3.0, help="Rainbow spread" - ) - parser.add_option( - "-F", "--freq", type="float", default=0.1, help="Rainbow frequency" - ) - parser.add_option("-S", "--seed", type="int", default=0, help="Rainbow seed") - parser.add_option( - "-a", - "--animate", - action="store_true", - default=False, - help="Enable psychedelics", - ) - parser.add_option( - "-d", "--duration", type="int", default=12, help="Animation duration" - ) - parser.add_option( - "-s", "--speed", type="float", default=20.0, help="Animation speed" - ) - parser.add_option( - "-f", - "--force", - action="store_true", - default=False, - help="Force colour even when stdout is not a tty", - ) - - parser.add_option( - "-3", action="store_const", dest="mode", const=8, help="Force 3 bit colour mode" - ) - parser.add_option( - "-4", - action="store_const", - dest="mode", - const=16, - help="Force 4 bit colour mode", - ) - parser.add_option( - "-8", - action="store_const", - dest="mode", - const=256, - help="Force 8 bit colour mode", - ) - - parser.add_option( - "-c", - "--charset-py2", - default="utf-8", - help="Manually set a charset to convert from, for python 2.7", - ) - - options, args = parser.parse_args() - options.os = random.randint(0, 256) if options.seed == 0 else options.seed - options.mode = options.mode or detect_mode() - - if os.name == "nt": - lolcat = LolCat(mode=options.mode, output=stdoutWin()) - else: - lolcat = LolCat(mode=options.mode) - - if not args: - args = ["-"] - - for filename in args: - try: - if filename == "-": - lolcat.cat(sys.stdin, options) - else: - with open(filename, "r", errors="backslashreplace") as handle: - lolcat.cat(handle, options) - except IOError as error: - sys.stderr.write(str(error) + "\n") - except KeyboardInterrupt: - sys.stderr.write("\n") - # exit 130 for terminated-by-ctrl-c, from http://tldp.org/LDP/abs/html/exitcodes.html - return 130 - - -if __name__ == "__main__": - sys.exit(run())