Fix the --profile and --region flags, fix logging, and add support fo…

…r the AWS_SESSION_TOKEN environment variable

CHANGELOG.md

@@ -9,6 +9,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 Nothing currently!
 
+## v0.4.0 - 2023-10-18
+
+### Added
+
+* Support for an optional `AWS_SESSION_TOKEN` environment variable is now implemented (by
+  [@cooperwalbrun](https://github.com/cooperwalbrun))
+
+### Fixed
+
+* The `--profile` and `--region` CLI arguments now function as expected (by
+  [@cooperwalbrun](https://github.com/cooperwalbrun))
+* All CIDRs for VPCs which were omitted during the handling of the `--prefix` flag are now logged
+  as expected (by [@cooperwalbrun](https://github.com/cooperwalbrun))
+
 ## v0.3.2 - 2022-11-20
 
 ### Changed

README.md

@@ -123,8 +123,9 @@ a keypair. The former may be specified using the `--profile` argument on the CLI
 must be specified in environment variables. If both are available simultaneously, `aws-cidr-finder`
 will prefer the profile.
 
-The environment variables for the keypair approach are `AWS_ACCESS_KEY_ID` and
-`AWS_SECRET_ACCESS_KEY` (the same values Boto uses).
+The environment variables for the keypair approach are `AWS_ACCESS_KEY_ID`,
+`AWS_SECRET_ACCESS_KEY`, and optionally `AWS_SESSION_TOKEN` (if authenticating with a session).
+These are the same values Boto uses.
 
 You should also ensure that the profile/keypair you are using has the AWS IAM access needed to make
 the underlying API calls via Boto. Here is a minimal IAM policy document that fills this

src/aws_cidr_finder/__main__.py

@@ -1,4 +1,5 @@
 import json
+import os
 import sys
 from argparse import ArgumentParser, Namespace
 from typing import Any
@@ -57,7 +58,16 @@ def _parse_arguments(arguments: list[str]) -> dict[str, Any]:
 def main() -> None:
     arguments = _parse_arguments(_get_arguments())
 
-    boto = BotoWrapper(profile_name=arguments.get("PROFILE"), region=arguments.get("REGION"))
+    if arguments.get("profile") is None and (os.environ.get("AWS_ACCESS_KEY_ID") is None
+                                             or os.environ.get("AWS_SECRET_ACCESS_KEY")):
+        print((
+            "You must specify either a profile or an access keypair via the environment variables "
+            "AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and optionally AWS_SESSION_TOKEN (if "
+            "authenticating with a session)"
+        ))
+        exit(1)
+
+    boto = BotoWrapper(profile_name=arguments.get("profile"), region=arguments.get("region"))
 
     ipv6: bool = arguments["ipv6"]
 
@@ -76,7 +86,7 @@ def main() -> None:
                 subnet_cidr_gaps[vpc], arguments["prefix"]
             )
             subnet_cidr_gaps[vpc] = converted_cidrs
-            messages = m
+            messages += m
 
     if arguments["json"]:
         output: JSONOutput = {"aws-cidr-finder-messages": messages, "vpcs": {}}

src/aws_cidr_finder/boto_wrapper.py

@@ -55,6 +55,7 @@ def __init__(self, *, profile_name: Optional[str], region: Optional[str]):
             boto = boto3.session.Session(
                 aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
                 aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
+                aws_session_token=os.environ.get("AWS_SESSION_TOKEN"),
                 region_name=region
             )
         self._client: EC2Client = boto.client("ec2")

tests/test___main__.py

@@ -1,6 +1,7 @@
 import json
 from unittest.mock import call, MagicMock
 
+import pytest
 from pytest_mock import MockerFixture
 from tabulate import tabulate
 
@@ -9,13 +10,30 @@
 
 
 def test_main_no_arguments(mocker: MockerFixture) -> None:
+    mocker.patch("aws_cidr_finder.__main__._get_arguments", return_value=[])
+    print_mock: MagicMock = mocker.patch("builtins.print")
+
+    with pytest.raises(SystemExit) as wrapped_exit_code:
+        __main__.main()
+
+    assert wrapped_exit_code.value.code == 1
+    print_mock.assert_has_calls([
+        call((
+            "You must specify either a profile or an access keypair via the environment variables "
+            "AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and optionally AWS_SESSION_TOKEN (if "
+            "authenticating with a session)"
+        ))
+    ])
+
+
+def test_main_only_profile_argument(mocker: MockerFixture) -> None:
     boto_wrapper_mock = MagicMock()
     boto_wrapper_mock.get_vpc_data = lambda ipv6: [
         VPC(id="test", name="test-vpc", 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", return_value=boto_wrapper_mock)
-    mocker.patch("aws_cidr_finder.__main__._get_arguments", return_value=[])
+    mocker.patch("aws_cidr_finder.__main__._get_arguments", return_value=["--profile", "test"])
     print_mock: MagicMock = mocker.patch("builtins.print")
 
     __main__.main()
@@ -37,7 +55,8 @@ def test_main_json_output(mocker: MockerFixture) -> None:
     mocker.patch("aws_cidr_finder.__main__.BotoWrapper", return_value=boto_wrapper_mock)
     mocker.patch("aws_cidr_finder.__main__.BotoWrapper", return_value=boto_wrapper_mock)
     mocker.patch(
-        "aws_cidr_finder.__main__._get_arguments", return_value=["--json", "--prefix", "20"]
+        "aws_cidr_finder.__main__._get_arguments",
+        return_value=["--profile", "test", "--json", "--prefix", "20"]
     )
     print_mock: MagicMock = mocker.patch("builtins.print")