Merge pull request #24 from cooperwalbrun/python-api

Refactor the Codebase to Support Python-Based Use Cases

.coveragerc

@@ -1,8 +1,6 @@
 [run]
 branch = true
 source = aws_cidr_finder
-omit =
-    **/aws_cidr_finder/__init__.py
 
 [report]
 fail_under = 90

.github/workflows/master.yml

@@ -13,9 +13,9 @@ jobs:
         python-version: ["3.10", "3.11"]
     steps:
       - name: Checkout source code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies

.github/workflows/pull-requests.yml

@@ -14,9 +14,9 @@ jobs:
         python-version: ["3.10", "3.11"]
     steps:
       - name: Checkout source code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies

.github/workflows/tags.yml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout source code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Set up Python 3.11
         uses: actions/setup-python@v4
         with:

CHANGELOG.md

@@ -9,6 +9,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 Nothing currently!
 
+## v0.5.0 - 2024-01-17
+
+### Added
+
+* Added a Python-based API for using `aws-cidr-finder` in a programmatic fashion (by
+  [@cooperwalbrun](https://github.com/cooperwalbrun))
+
+### Changed
+
+* The JSON format of `aws-cidr-finder`'s output is now simpler: the `aws-cidr-finder-messages` key
+  has been changed to `messages`, `vpcs` has been changed to `data`, and the structure of `data` is
+  now a flat list of "VPC" `dict`s where each `dict` has the following keys: `id`, `name`, `cidr`,
+  and `available_cidr_blocks` (by [@cooperwalbrun](https://github.com/cooperwalbrun))
+
 ## v0.4.1 - 2023-10-18
 
 ### Changed

README.md

@@ -4,7 +4,10 @@
    1. [An Example](#an-example)
 2. [Installation](#installation)
 3. [Configuration](#configuration)
-4. [Contributing](#contributing)
+4. [Usage](#usage)
+   1. [CLI](#cli)
+   2. [Python](#python)
+5. [Contributing](#contributing)
 
 ## Overview
 
@@ -145,6 +148,54 @@ requirement:
 Read more about the actions shown above
 [here](https://docs.aws.amazon.com/service-authorization/latest/reference/list_amazonec2.html).
 
+## Usage
+
+### CLI
+
+See [An Example](#an-example) above for a detailed demonstration of the CLI interface of this tool.
+You can also use `aws-cidr-finder --help` to see command line options.
+
+### Python
+
+Setup:
+
+```python
+from aws_cidr_finder import JSONOutput, find_cidrs
+
+# All arguments
+output: JSONOutput = find_cidrs(profile_name="", region="", ipv6=False, desired_prefix=20)
+
+# Minimal arguments (profile-based authentication)
+output: JSONOutput = find_cidrs(profile_name="")
+
+# Minimal arguments (environment variable-based authentication)
+output: JSONOutput = find_cidrs()
+
+# Other miscellaneous combinations
+output: JSONOutput = find_cidrs(profile_name="", ipv6=True)
+output: JSONOutput = find_cidrs(profile_name="", desired_prefix=16)
+output: JSONOutput = find_cidrs(region="")
+# ...and so on
+```
+
+Accessing the CIDR data:
+
+```python
+output: JSONOutput = ...  # See above
+
+for message in output["messages"]:
+    # Print the messages that would have been written to STDOUT when using the CLI
+    print(message)
+
+for vpc in output["data"]:
+    # The following shows how to access all available fields in the data object
+    print(vpc["id"])
+    print(vpc["name"])
+    print(vpc["cidr"])
+    for cidr in vpc["available_cidr_blocks"]:
+        print(cidr)
+```
+
 ## Contributing
 
 See [CONTRIBUTING.md](CONTRIBUTING.md) for developer-oriented information.

src/aws_cidr_finder/__init__.py

@@ -1,5 +1,11 @@
+from typing import Optional
+
 from importlib_metadata import PackageNotFoundError, version
 
+from aws_cidr_finder import custom_types
+from aws_cidr_finder.boto_wrapper import BotoWrapper
+from aws_cidr_finder.core import convert_to_json_format
+
 try:
     # We hard-code the name rather than using __name__ because the package name has an underscore
     # instead of a hyphen
@@ -9,3 +15,17 @@
     __version__ = "unknown"
 finally:
     del version, PackageNotFoundError
+
+JSONOutput = custom_types.JSONOutput
+
+
+def find_cidrs(
+    *,
+    profile_name: Optional[str] = None,
+    region: Optional[str] = None,
+    ipv6: bool = False,
+    desired_prefix: Optional[int] = None
+) -> JSONOutput:
+    boto = BotoWrapper(profile_name=profile_name, region=region)
+    subnet_cidr_gaps, messages = boto.get_subnet_cidr_gaps(ipv6=ipv6, prefix=desired_prefix)
+    return convert_to_json_format(subnet_cidr_gaps, messages)

src/aws_cidr_finder/__main__.py

@@ -8,7 +8,7 @@
 
 from aws_cidr_finder import core
 from aws_cidr_finder.boto_wrapper import BotoWrapper
-from aws_cidr_finder.custom_types import SingleCIDRVPC, JSONOutput
+from aws_cidr_finder.core import convert_to_json_format
 
 _parser: ArgumentParser = ArgumentParser(
     description="A CLI tool for finding unused CIDR blocks in AWS VPCs."
@@ -71,30 +71,12 @@ def main() -> None:
 
     ipv6: bool = arguments["ipv6"]
 
-    subnet_cidr_gaps: dict[SingleCIDRVPC, list[str]] = {}
-    messages: list[str] = []
-
-    for vpc in core.split_out_individual_cidrs(boto.get_vpc_data(ipv6=ipv6)):
-        # yapf: disable
-        subnet_cidr_gaps[vpc] = core.find_subnet_holes(
-            vpc.cidr,
-            vpc.subnets
-        )
-        # yapf: enable
-        if arguments.get("prefix") is not None:
-            converted_cidrs, m = core.break_down_to_desired_prefix(
-                vpc.readable_name, subnet_cidr_gaps[vpc], arguments["prefix"]
-            )
-            subnet_cidr_gaps[vpc] = converted_cidrs
-            messages += m
+    subnet_cidr_gaps, messages = boto.get_subnet_cidr_gaps(
+        ipv6=ipv6, prefix=arguments.get("prefix")
+    )
 
     if arguments["json"]:
-        output: JSONOutput = {"aws-cidr-finder-messages": messages, "vpcs": {}}
-        for vpc, subnet_cidrs in subnet_cidr_gaps.items():
-            if vpc.readable_name not in output["vpcs"]:
-                output["vpcs"][vpc.readable_name] = {}
-            output["vpcs"][vpc.readable_name][vpc.cidr] = core.sort_cidrs(subnet_cidrs)
-        print(json.dumps(output))
+        print(json.dumps(convert_to_json_format(subnet_cidr_gaps, messages)))
     else:
         if len(subnet_cidr_gaps) == 0:
             print(f"No available {'IPv6' if ipv6 else 'IPv4'} CIDR blocks were found in any VPC.")

src/aws_cidr_finder/boto_wrapper.py

@@ -6,7 +6,8 @@
 from mypy_boto3_ec2.type_defs import TagTypeDef, VpcTypeDef, DescribeSubnetsResultTypeDef, \
     SubnetTypeDef
 
-from aws_cidr_finder.custom_types import VPC
+from aws_cidr_finder import core
+from aws_cidr_finder.custom_types import VPC, SingleCIDRVPC
 
 
 def _get_vpc_name(tags: list[TagTypeDef]) -> Optional[str]:
@@ -47,8 +48,8 @@ def _parse_subnet_cidrs(subnets: list[SubnetTypeDef], *, ipv6: bool) -> list[str
         return [subnet["CidrBlock"] for subnet in subnets if "CidrBlock" in subnet]
 
 
-class BotoWrapper:  # pragma: no cover
-    def __init__(self, *, profile_name: Optional[str], region: Optional[str]):
+class BotoWrapper:
+    def __init__(self, *, profile_name: Optional[str], region: Optional[str]):  # pragma: no cover
         if profile_name is not None:
             boto = boto3.session.Session(profile_name=profile_name, region_name=region)
         else:
@@ -60,7 +61,7 @@ def __init__(self, *, profile_name: Optional[str], region: Optional[str]):
             )
         self._client: EC2Client = boto.client("ec2")
 
-    def get_vpc_data(self, *, ipv6: bool) -> list[VPC]:
+    def _get_vpc_data(self, *, ipv6: bool) -> list[VPC]:  # pragma: no cover
         vpcs = self._client.describe_vpcs()["Vpcs"]
         return [
             VPC(
@@ -73,5 +74,27 @@ def get_vpc_data(self, *, ipv6: bool) -> list[VPC]:
             ) for vpc in vpcs
         ]
 
-    def _get_subnet_cidrs(self, vpc_id: str) -> DescribeSubnetsResultTypeDef:
+    def _get_subnet_cidrs(self, vpc_id: str) -> DescribeSubnetsResultTypeDef:  # pragma: no cover
         return self._client.describe_subnets(Filters=[{"Name": "vpc-id", "Values": [vpc_id]}])
+
+    def get_subnet_cidr_gaps(
+        self, *, ipv6: bool, prefix: Optional[int]
+    ) -> tuple[dict[SingleCIDRVPC, list[str]], list[str]]:
+        subnet_cidr_gaps: dict[SingleCIDRVPC, list[str]] = {}
+        messages: list[str] = []
+
+        for vpc in core.split_out_individual_cidrs(self._get_vpc_data(ipv6=ipv6)):
+            # yapf: disable
+            subnet_cidr_gaps[vpc] = core.find_subnet_holes(
+                vpc.cidr,
+                vpc.subnets
+            )
+            # yapf: enable
+            if prefix is not None:
+                converted_cidrs, m = core.break_down_to_desired_prefix(
+                    vpc.readable_name, subnet_cidr_gaps[vpc], prefix
+                )
+                subnet_cidr_gaps[vpc] = converted_cidrs
+                messages += m
+
+        return subnet_cidr_gaps, messages

src/aws_cidr_finder/core.py

@@ -3,7 +3,7 @@
     IPv4Address, IPv6Address
 from typing import Optional, Union
 
-from aws_cidr_finder.custom_types import VPC, SingleCIDRVPC
+from aws_cidr_finder.custom_types import VPC, SingleCIDRVPC, JSONOutput, VPCCIDRData
 
 
 def _get_cidr(network: Union[IPv4Network, IPv6Network]) -> str:
@@ -167,3 +167,18 @@ def break_down_to_desired_prefix(readable_vpc_name: str, cidrs: list[str],
             ret.append(_get_cidr(sub))
 
     return ret, messages
+
+
+def convert_to_json_format(
+    subnet_cidr_gaps: dict[SingleCIDRVPC, list[str]], messages: list[str]
+) -> JSONOutput:
+    # yapf: disable
+    vpc_data: list[VPCCIDRData] = [{
+        "id": vpc.id,
+        "name": vpc.name,
+        "cidr": vpc.cidr,
+        "available_cidr_blocks": subnet_cidrs
+    } for vpc, subnet_cidrs in subnet_cidr_gaps.items()]
+    # yapf: enable
+
+    return {"messages": messages, "data": vpc_data}

src/aws_cidr_finder/custom_types.py

@@ -36,7 +36,9 @@ def __hash__(self) -> int:
         return hash((self.id, self.cidr))
 
 
-JSONOutput = TypedDict("JSONOutput", {
-    "aws-cidr-finder-messages": list[str],
-    "vpcs": dict[str, dict[str, list[str]]]
-})
+VPCCIDRData = TypedDict(
+    "VPCCIDRData", {
+        "id": str, "name": Optional[str], "cidr": str, "available_cidr_blocks": list[str]
+    }
+)
+JSONOutput = TypedDict("JSONOutput", {"messages": list[str], "data": list[VPCCIDRData]})

tests/test___main__.py

@@ -27,11 +27,13 @@ def test_main_no_arguments(mocker: MockerFixture) -> None:
 
 
 def test_main_only_profile_argument(mocker: MockerFixture) -> None:
-    boto_wrapper_mock = MagicMock()
-    boto_wrapper_mock.get_vpc_data = lambda ipv6: [
-        VPC(id="test1", name="test-vpc1", cidrs=["172.31.0.0/19"], subnets=["172.31.0.0/20"])
-    ]
-    mocker.patch("aws_cidr_finder.__main__.BotoWrapper", return_value=boto_wrapper_mock)
+    mocker.patch("aws_cidr_finder.__main__.BotoWrapper.__init__", return_value=None)
+    mocker.patch(
+        "aws_cidr_finder.__main__.BotoWrapper._get_vpc_data",
+        return_value=[
+            VPC(id="test1", name="test-vpc1", cidrs=["172.31.0.0/19"], subnets=["172.31.0.0/20"])
+        ]
+    )
     mocker.patch("aws_cidr_finder.__main__._get_arguments", return_value=["--profile", "test"])
     print_mock: MagicMock = mocker.patch("builtins.print")
 
@@ -47,12 +49,14 @@ def test_main_only_profile_argument(mocker: MockerFixture) -> None:
 
 
 def test_main_json_output(mocker: MockerFixture) -> None:
-    boto_wrapper_mock = MagicMock()
-    boto_wrapper_mock.get_vpc_data = lambda ipv6: [
-        VPC(id="test1", name="test-vpc1", cidrs=["172.31.0.0/19"], subnets=["172.31.0.0/20"]),
-        VPC(id="test2", name="test-vpc2", cidrs=["172.31.32.0/20"], subnets=["172.31.32.0/21"])
-    ]
-    mocker.patch("aws_cidr_finder.__main__.BotoWrapper", return_value=boto_wrapper_mock)
+    mocker.patch("aws_cidr_finder.__main__.BotoWrapper.__init__", return_value=None)
+    mocker.patch(
+        "aws_cidr_finder.__main__.BotoWrapper._get_vpc_data",
+        return_value=[
+            VPC(id="test1", name="test-vpc1", cidrs=["172.31.0.0/19"], subnets=["172.31.0.0/20"]),
+            VPC(id="test2", name="test-vpc2", cidrs=["172.31.32.0/20"], subnets=["172.31.32.0/21"])
+        ]
+    )
     mocker.patch(
         "aws_cidr_finder.__main__._get_arguments",
         return_value=["--profile", "test", "--json", "--prefix", "20"]
@@ -61,21 +65,29 @@ def test_main_json_output(mocker: MockerFixture) -> None:
 
     __main__.main()
 
+    # yapf: disable
     print_mock.assert_has_calls([
         call(
             json.dumps({
-                "aws-cidr-finder-messages": [(
+                "messages": [(
                     "Note: skipping the CIDR '172.31.40.0/21' in the VPC 'test-vpc2' because its "
                     "prefix (21) is numerically greater than the requested prefix (20)"
                 )],
-                "vpcs": {
-                    "test-vpc1": {
-                        "172.31.0.0/19": ["172.31.16.0/20"]
+                "data": [
+                    {
+                        "id": "test1",
+                        "name": "test-vpc1",
+                        "cidr": "172.31.0.0/19",
+                        "available_cidr_blocks": ["172.31.16.0/20"]
                     },
-                    "test-vpc2": {
-                        "172.31.32.0/20": []
+                    {
+                        "id": "test2",
+                        "name": "test-vpc2",
+                        "cidr": "172.31.32.0/20",
+                        "available_cidr_blocks": []
                     }
-                }
+                ]
             })
         )
     ])
+    # yapf: enable