diff --git a/README.md b/README.md index 2117be8..5fbc41e 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ pip install digaws ## Usage ```text -usage: digaws [-h] [--output ] [--debug] [ ...] +usage: digaws [-h] [--output ] [--output-fields [{prefix,region,service,network_border_group} ...]] [--debug] [--version] [ ...] Look up canonical information for AWS IP addresses and networks @@ -30,7 +30,10 @@ optional arguments: -h, --help show this help message and exit --output Formatting style for command output, by default plain + --output-fields [{prefix,region,service,network_border_group} ...] + Print only the given fields --debug Enable debug + --version show program's version number and exit ``` ## Examples @@ -104,3 +107,14 @@ Network border group: sa-east-1 } ] ``` + +- Choose output fields + +```text +~ ยป digaws 34.255.166.63 --output-fields service + +Service: AMAZON + +Service: EC2 + +``` diff --git a/digaws/digaws.py b/digaws/digaws.py index 2ebb6c8..8ed391f 100755 --- a/digaws/digaws.py +++ b/digaws/digaws.py @@ -20,6 +20,7 @@ from . import __version__ AWS_IP_RANGES_URL = 'https://ip-ranges.amazonaws.com/ip-ranges.json' CACHE_DIR = Path(Path.home() / '.digaws') CACHE_FILE = CACHE_DIR / 'ip-ranges.json' +OUTPUT_FIELDS = ['prefix', 'region', 'service', 'network_border_group'] logger = logging.getLogger() handler = logging.StreamHandler(sys.stderr) @@ -87,19 +88,23 @@ class UnexpectedRequestException(Exception): class DigAWSPrettyPrinter: - def __init__(self, data: List[Dict], output_filter: List[str] = []): + def __init__(self, data: List[Dict], output_fields: List[str] = []): self.data = data - self.output_filter = output_filter + self.output_fields = output_fields def plain_print(self) -> None: for prefix in self.data: - try: - print(f'Prefix: {prefix["ip_prefix"]}') - except KeyError: - print(f'IPv6 Prefix: {prefix["ipv6_prefix"]}') - print(f'Region: {prefix["region"]}') - print(f'Service: {prefix["service"]}') - print(f'Network border group: {prefix["network_border_group"]}') + if 'prefix' in self.output_fields: + try: + print(f'Prefix: {prefix["ip_prefix"]}') + except KeyError: + print(f'IPv6 Prefix: {prefix["ipv6_prefix"]}') + if 'region' in self.output_fields: + print(f'Region: {prefix["region"]}') + if 'service' in self.output_fields: + print(f'Service: {prefix["service"]}') + if 'network_border_group' in self.output_fields: + print(f'Network border group: {prefix["network_border_group"]}') print('') def json_print(self) -> None: @@ -110,22 +115,26 @@ class DigAWSPrettyPrinter: prefix_type = 'ip_prefix' except KeyError: prefix_type = 'ipv6_prefix' - data.append( - { - prefix_type: str(prefix[prefix_type]), - 'region': prefix['region'], - 'service': prefix['service'], - 'network_border_group': prefix['network_border_group'] - } - ) + + item_dict = {} + if 'prefix' in self.output_fields: + item_dict.update({prefix_type: str(prefix[prefix_type])}) + if 'region' in self.output_fields: + item_dict.update({'region': prefix['region']}) + if 'service' in self.output_fields: + item_dict.update({'service': prefix['service']}) + if 'network_border_group' in self.output_fields: + item_dict.update({'network_border_group': prefix['network_border_group']}) + data.append(item_dict) + if data: print(json.dumps(data, indent=2)) class DigAWS(): - def __init__(self, ip_ranges: Dict, output: str = 'plain', output_filter: List[str] = []): + def __init__(self, *, ip_ranges: Dict, output: str = 'plain', output_fields: List[str] = []): self.output = output - self.output_filter = output_filter + self.output_fields = output_fields self.ip_prefixes = [ { 'ip_prefix': ipaddress.IPv4Network(prefix['ip_prefix']), @@ -148,7 +157,7 @@ class DigAWS(): def lookup(self, address: str) -> DigAWSPrettyPrinter: return DigAWSPrettyPrinter( self._lookup_data(address), - self.output_filter + self.output_fields ) def _lookup_data(self, address: str) -> List[Dict]: @@ -193,6 +202,15 @@ def arguments_parser() -> argparse.ArgumentParser: default='plain', help='Formatting style for command output, by default %(default)s' ) + parser.add_argument( + '--output-fields', + nargs='*', + choices=OUTPUT_FIELDS, + required=False, + dest='output_fields', + default=OUTPUT_FIELDS, + help='Print only the given fields' + ) parser.add_argument( '--debug', action='store_true', @@ -225,7 +243,7 @@ def main(): try: ip_ranges = get_aws_ip_ranges() - dig = DigAWS(ip_ranges) + dig = DigAWS(ip_ranges=ip_ranges, output_fields=args.output_fields) responses = [] for address in args.addresses: @@ -242,7 +260,7 @@ def main(): for response in responses: joined += response.data - DigAWSPrettyPrinter(joined).json_print() + DigAWSPrettyPrinter(joined, args.output_fields).json_print() except ( RequestException, ipaddress.AddressValueError, diff --git a/tests/__init__.py b/tests/__init__.py index a31b609..5485a04 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -111,6 +111,18 @@ RESPONSE_JSON_PRINT = '''[ ] ''' +RESPONSE_JSON_FIELDS_PRINT = '''[ + { + "service": "ROUTE53_HEALTHCHECKS", + "network_border_group": "us-west-2" + }, + { + "service": "EC2", + "network_border_group": "us-west-2" + } +] +''' + RESPONSE_JSON_JOINED_PRINT = '''[ { "ip_prefix": "52.94.76.0/22", diff --git a/tests/test_dig_aws.py b/tests/test_dig_aws.py index 9eb82a1..8bf9588 100644 --- a/tests/test_dig_aws.py +++ b/tests/test_dig_aws.py @@ -11,7 +11,10 @@ import tests @pytest.fixture def test_dig(): - return digaws.DigAWS(json.loads(tests.AWS_IP_RANGES)) + return digaws.DigAWS(ip_ranges=json.loads( + tests.AWS_IP_RANGES), + output_fields=digaws.OUTPUT_FIELDS + ) def test_cli(capsys): @@ -44,6 +47,25 @@ def test_cli_invocation(capsys, mocker): assert out == tests.RESPONSE_JSON_JOINED_PRINT +def test_cli_output_plain_fields_invocation(capsys, mocker): + sys.argv = ['digaws', '52.94.76.0/22', '--output=plain', '--output-fields', 'region'] + mocker.patch('digaws.digaws.get_aws_ip_ranges', return_value=json.loads(tests.AWS_IP_RANGES)) + digaws.main() + out, _ = capsys.readouterr() + + assert out == 'Region: us-west-2\n\n' + + +def test_cli_output_json_fields_invocation(capsys, mocker): + sys.argv = ['digaws', '2600:1f14:fff:f810:a1c1:f507:a2d1:2dd8', '--output=json', + '--output-fields', 'service', 'network_border_group'] + mocker.patch('digaws.digaws.get_aws_ip_ranges', return_value=json.loads(tests.AWS_IP_RANGES)) + digaws.main() + out, _ = capsys.readouterr() + + assert out == tests.RESPONSE_JSON_FIELDS_PRINT + + def test_dig_aws_construct(test_dig): assert test_dig.ip_prefixes == tests.AWS_IPV4_RANGES_OBJ assert test_dig.ipv6_prefixes == tests.AWS_IPV6_RANGES_OBJ