From c003cc81859f13cc699c4f7994103a69450f5c7d Mon Sep 17 00:00:00 2001 From: Cooper Walbrun <69944858+cooperwalbrun@users.noreply.github.com> Date: Wed, 18 Oct 2023 20:45:22 -0500 Subject: [PATCH] Fix the --profile and --region flags, fix logging, and add support for the AWS_SESSION_TOKEN environment variable --- CHANGELOG.md | 14 ++++++++++++++ README.md | 5 +++-- src/aws_cidr_finder/__main__.py | 14 ++++++++++++-- src/aws_cidr_finder/boto_wrapper.py | 1 + tests/test___main__.py | 23 +++++++++++++++++++++-- 5 files changed, 51 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 40f85d3..fc75e72 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index efbef62..735ab1c 100644 --- a/README.md +++ b/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 diff --git a/src/aws_cidr_finder/__main__.py b/src/aws_cidr_finder/__main__.py index 29f41ff..5ca9c12 100644 --- a/src/aws_cidr_finder/__main__.py +++ b/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": {}} diff --git a/src/aws_cidr_finder/boto_wrapper.py b/src/aws_cidr_finder/boto_wrapper.py index b8c8827..6ccacd9 100644 --- a/src/aws_cidr_finder/boto_wrapper.py +++ b/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") diff --git a/tests/test___main__.py b/tests/test___main__.py index e798ba3..cdefbd7 100644 --- a/tests/test___main__.py +++ b/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")