Merge pull request #6 from dcarrillo/1.0.4

Return an empty array when there is no data for json output

.github/workflows/main.yml

@@ -4,6 +4,8 @@ on:
   push:
     branches:
       - main
+    tags:
+      - '*'
   pull_request:
 
 jobs:
@@ -36,7 +38,8 @@ jobs:
 
   build_publish:
     runs-on: ubuntu-latest
-    if: github.event_name == 'push'
+    needs: tests
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
     steps:
       - uses: actions/checkout@v2
 
@@ -48,7 +51,7 @@ jobs:
       - name: Install tools
         run: |
           python -m pip install --upgrade pip
-          pip install twine
+          pip install twine wheel
 
       - name: Build
         run: |

README.md

@@ -4,6 +4,9 @@ The digaws lookup tool displays information for a given IP address (v4 o v6) or
 In order to save bandwidth and time this tool requests the [AWS IP ranges](https://ip-ranges.amazonaws.com/ip-ranges.json) and keeps
 a cached version until a new version is published.
 
+![PyPI version](https://badge.fury.io/py/digaws.svg)
+![CI](https://github.com/dcarrillo/digaws/workflows/CI/badge.svg)
+
 ## Requirements
 
 Python >= 3.7
@@ -19,7 +22,7 @@ pip install digaws
 ## Usage
 
 ```text
-usage: digaws [-h] [--output <plain|json>] [--debug] <ip address|cidr> [<ip address|cidr> ...]
+usage: digaws [-h] [--output <plain|json>] [--output-fields [{prefix,region,service,network_border_group} ...]] [--debug] [--version] <ip address|cidr> [<ip address|cidr> ...]
 
 Look up canonical information for AWS IP addresses and networks
 
@@ -30,7 +33,10 @@ optional arguments:
   -h, --help            show this help message and exit
   --output <plain|json>
                         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 +110,14 @@ Network border group: sa-east-1
   }
 ]
 ```
+
+- Choose output fields
+
+```text
+~ » digaws 34.255.166.63 --output-fields service
+
+Service: AMAZON
+
+Service: EC2
+
+```

digaws/__init__.py

@@ -1,2 +1,2 @@
-__version__ = '1.0.1'
+__version__ = '1.0.4'
 __description__ = 'Look up canonical information for AWS IP addresses and networks'

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,25 @@ 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']
-                }
-            )
-        if data:
-            print(json.dumps(data, indent=2))
+
+            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)
+
+        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 +156,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 +201,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 +242,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 +259,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,

setup.py

@@ -17,7 +17,7 @@ setup(
     url='http://github.com/dcarrillo/digaws',
     author='Daniel Carrillo',
     author_email='daniel.carrillo@gmail.com',
-    license='MIT',
+    license='Apache Software License',
     packages=['digaws'],
     zip_safe=False,
     classifiers=[

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",

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

tests/test_get_aws_ip_ranges_json.py

@@ -84,4 +84,5 @@ def test_get_aws_ip_ranges_invalid_status(mocker, fs, create_cache_dir) -> None:
 
     with pytest.raises(digaws.UnexpectedRequestException) as e:
         digaws.get_aws_ip_ranges()
-        assert e.startswith('Unexpected response from')
+
+    assert e.match('^Unexpected response from')