Skip to content

Commit c105a15

Browse files
RichieRunnerRichieRunner
authored
add EC2 to allow/revoke security group ingress (#72)
Signed-off-by: RichieRunner <richie.ho@outlook.com>
1 parent d506b21 commit c105a15

File tree

4 files changed

+261
-5
lines changed

4 files changed

+261
-5
lines changed

‎CHANGELOG.md

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,15 @@
22

33
## [Unreleased][]
44

5-
[Unreleased]: https://github.com/chaostoolkit-incubator/chaostoolkit-aws/compare/0.13.0...HEAD
5+
[Unreleased]: https://github.com/chaostoolkit-incubator/chaostoolkit-aws/compare/0.14.0...HEAD
6+
7+
## [0.14.0][]
8+
9+
[0.14.0]: https://github.com/chaostoolkit-incubator/chaostoolkit-aws/compare/0.13.0...0.14.0
10+
11+
### Added
12+
13+
- add EC2 actions to allow/revoke security group ingress
614

715
## [0.13.0][]
816

‎README.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Chaos Toolkit Extension for AWS
1+
# [Chaos Toolkit Extension for AWS](https://docs.chaostoolkit.org/drivers/aws/)
22

33
[![Build Status](https://travis-ci.org/chaostoolkit-incubator/chaostoolkit-aws.svg?branch=master)](https://travis-ci.org/chaostoolkit-incubator/chaostoolkit-aws)
44
[![Python versions](https://img.shields.io/pypi/pyversions/chaostoolkit-aws.svg)](https://www.python.org/)
@@ -368,7 +368,7 @@ def test_reboot_instance(aws_client):
368368
```
369369

370370
By using the [built-in Python module to mock objects][pymock], we can mock the
371-
EC2 client and assert we edo indeed call the appropriate method with the right
371+
EC2 client and assert that we do indeed call the appropriate method with the right
372372
arguments. You are encouraged to write more than a single test for various
373373
conditions.
374374

‎chaosaws/ec2/actions.py

Lines changed: 108 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from botocore.exceptions import ClientError
99
from chaosaws import aws_client
1010
from chaosaws.types import AWSResponse
11-
from chaoslib.exceptions import FailedActivity
11+
from chaoslib.exceptions import FailedActivity, ActivityFailed
1212
from chaoslib.types import Configuration, Secrets
1313
from logzero import logger
1414

@@ -657,3 +657,110 @@ def attach_instance_volume(client: boto3.client,
657657
'Unable to attach volume %s to instance %s: %s' % (
658658
volume_id, instance_id, e.response['Error']['Message']))
659659
return response
660+
661+
662+
def authorize_security_group_ingress(requested_security_group_id: str,
663+
ip_protocol: str,
664+
from_port: int,
665+
to_port: int,
666+
ingress_security_group_id: str = None,
667+
cidr_ip: str = None,
668+
configuration: Configuration = None,
669+
secrets: Secrets = None) -> AWSResponse:
670+
"""
671+
Add one ingress rule to a security group
672+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#EC2.Client.authorize_security_group_ingress
673+
674+
- requested_security_group_id: the id for the security group to update
675+
- ip_protocol: ip protocol name (tcp, udp, icmp, icmpv6) or -1 to specify all
676+
- from_port: start of port range
677+
- to_port: end of port range
678+
- ingress_security_group_id: id of the securiy group to allow access to. You can either specify this or cidr_ip.
679+
- cidr_ip: the IPv6 CIDR range.
680+
You can either specify this or ingress_security_group_id
681+
""" # noqa: E501
682+
client = aws_client('ec2', configuration, secrets)
683+
request_kwargs = create_ingress_kwargs(
684+
requested_security_group_id,
685+
ip_protocol,
686+
from_port,
687+
to_port,
688+
ingress_security_group_id,
689+
cidr_ip
690+
)
691+
try:
692+
response = client.authorize_security_group_ingress(**request_kwargs)
693+
return response
694+
except ClientError as e:
695+
raise ActivityFailed(
696+
'Failed to add ingress rule: {}'.format(
697+
e.response["Error"]["Message"]))
698+
699+
700+
def revoke_security_group_ingress(requested_security_group_id: str,
701+
ip_protocol: str,
702+
from_port: int,
703+
to_port: int,
704+
ingress_security_group_id: str = None,
705+
cidr_ip: str = None,
706+
configuration: Configuration = None,
707+
secrets: Secrets = None) -> AWSResponse:
708+
"""
709+
Remove one ingress rule from a security group
710+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#EC2.Client.revoke_security_group_ingress
711+
712+
- requested_security_group_id: the id for the security group to update
713+
- ip_protocol: ip protocol name (tcp, udp, icmp, icmpv6) or -1 to specify all
714+
- from_port: start of port range
715+
- to_port: end of port range
716+
- ingress_security_group_id: id of the securiy group to allow access to. You can either specify this or cidr_ip.
717+
- cidr_ip: the IPv6 CIDR range. You can either specify this or ingress_security_group_id
718+
""" # noqa: E501
719+
client = aws_client('ec2', configuration, secrets)
720+
request_kwargs = create_ingress_kwargs(
721+
requested_security_group_id,
722+
ip_protocol,
723+
from_port,
724+
to_port,
725+
ingress_security_group_id,
726+
cidr_ip
727+
)
728+
try:
729+
response = client.revoke_security_group_ingress(**request_kwargs)
730+
return response
731+
except ClientError as e:
732+
raise ActivityFailed(
733+
'Failed to remove ingress rule: {}'.format(
734+
e.response["Error"]["Message"]))
735+
736+
737+
def create_ingress_kwargs(requested_security_group_id: str,
738+
ip_protocol: str,
739+
from_port: int,
740+
to_port: int,
741+
ingress_security_group_id: str = None,
742+
cidr_ip: str = None,) -> Dict[str, any]:
743+
request_kwargs = {
744+
'GroupId': requested_security_group_id,
745+
'IpPermissions': [
746+
{
747+
'IpProtocol': ip_protocol,
748+
'IpRanges': [{
749+
# conditionally assign the following
750+
# 'CidrIp': cidr_ip
751+
}],
752+
'FromPort': from_port,
753+
'ToPort': to_port,
754+
'UserIdGroupPairs': [{
755+
# conditionally assign the following
756+
# 'GroupId': ingress_security_group_id
757+
}]
758+
}
759+
]
760+
}
761+
req = request_kwargs['IpPermissions'][0]
762+
if cidr_ip is not None:
763+
req['IpRanges'][0]['CidrIp'] = cidr_ip
764+
if ingress_security_group_id is not None:
765+
req['UserIdGroupPairs'][0]['GroupId'] = ingress_security_group_id
766+
return request_kwargs

‎tests/ec2/test_ec2_actions.py

Lines changed: 142 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@
66

77
from chaosaws.ec2.actions import (
88
stop_instance, stop_instances, terminate_instance, terminate_instances,
9-
start_instances, restart_instances, detach_random_volume, attach_volume)
9+
start_instances, restart_instances, detach_random_volume, attach_volume,
10+
authorize_security_group_ingress, revoke_security_group_ingress)
1011

1112
from chaoslib.exceptions import FailedActivity
1213

@@ -763,3 +764,143 @@ def test_attach_volume_ec2_filter(aws_client):
763764
InstanceId=instance_ids[0],
764765
VolumeId='vol-00000001')
765766
assert results[0]['DeviceName'] == '/dev/sdb'
767+
768+
769+
@patch('chaosaws.ec2.actions.aws_client', autospec=True)
770+
def test_authorize_security_group_ingress_with_ingress_security_group(aws_client):
771+
#arrange
772+
client = MagicMock()
773+
aws_client.return_value = client
774+
requested_security_group_id='sg-123456789abc'
775+
ingress_security_group_id='sg-123456789cba'
776+
ip_protocol='tcp'
777+
from_port=0
778+
to_port=80
779+
#act
780+
authorize_security_group_ingress(
781+
requested_security_group_id=requested_security_group_id,
782+
ip_protocol=ip_protocol,
783+
from_port=from_port,
784+
to_port=to_port,
785+
ingress_security_group_id=ingress_security_group_id
786+
)
787+
#assert
788+
client.authorize_security_group_ingress.assert_called_with(
789+
GroupId=requested_security_group_id,
790+
IpPermissions=[
791+
{
792+
'IpProtocol': ip_protocol,
793+
'FromPort': from_port,
794+
'ToPort': to_port,
795+
'IpRanges': [{}],
796+
'UserIdGroupPairs': [
797+
{
798+
'GroupId': ingress_security_group_id
799+
}
800+
]
801+
}
802+
]
803+
)
804+
805+
806+
@patch('chaosaws.ec2.actions.aws_client', autospec=True)
807+
def test_authorize_security_group_ingress_with_cidr_ip(aws_client):
808+
#arrange
809+
client = MagicMock()
810+
aws_client.return_value = client
811+
requested_security_group_id='sg-123456789abc'
812+
cidr_ip='0.0.0.0/0'
813+
ip_protocol='tcp'
814+
from_port=0
815+
to_port=80
816+
#act
817+
authorize_security_group_ingress(
818+
requested_security_group_id=requested_security_group_id,
819+
ip_protocol=ip_protocol,
820+
from_port=from_port,
821+
to_port=to_port,
822+
cidr_ip=cidr_ip
823+
)
824+
#assert
825+
client.authorize_security_group_ingress.assert_called_with(
826+
GroupId=requested_security_group_id,
827+
IpPermissions=[
828+
{
829+
'IpProtocol': ip_protocol,
830+
'FromPort': from_port,
831+
'ToPort': to_port,
832+
'IpRanges': [{'CidrIp': cidr_ip}],
833+
'UserIdGroupPairs': [{}]
834+
}
835+
]
836+
)
837+
838+
839+
@patch('chaosaws.ec2.actions.aws_client', autospec=True)
840+
def test_revoke_security_group_ingress_with_ingress_security_group(aws_client):
841+
#arrange
842+
client = MagicMock()
843+
aws_client.return_value = client
844+
requested_security_group_id='sg-123456789abc'
845+
ingress_security_group_id='sg-123456789cba'
846+
ip_protocol='tcp'
847+
from_port=0
848+
to_port=80
849+
#act
850+
revoke_security_group_ingress(
851+
requested_security_group_id=requested_security_group_id,
852+
ip_protocol=ip_protocol,
853+
from_port=from_port,
854+
to_port=to_port,
855+
ingress_security_group_id=ingress_security_group_id
856+
)
857+
#assert
858+
client.revoke_security_group_ingress.assert_called_with(
859+
GroupId=requested_security_group_id,
860+
IpPermissions=[
861+
{
862+
'IpProtocol': ip_protocol,
863+
'FromPort': from_port,
864+
'ToPort': to_port,
865+
'IpRanges': [{}],
866+
'UserIdGroupPairs': [
867+
{
868+
'GroupId': ingress_security_group_id
869+
}
870+
]
871+
}
872+
]
873+
)
874+
875+
876+
@patch('chaosaws.ec2.actions.aws_client', autospec=True)
877+
def test_revoke_security_group_ingress_with_cidr_ip(aws_client):
878+
#arrange
879+
client = MagicMock()
880+
aws_client.return_value = client
881+
requested_security_group_id='sg-123456789abc'
882+
cidr_ip='0.0.0.0/0'
883+
ip_protocol='tcp'
884+
from_port=0
885+
to_port=80
886+
#act
887+
revoke_security_group_ingress(
888+
requested_security_group_id=requested_security_group_id,
889+
ip_protocol=ip_protocol,
890+
from_port=from_port,
891+
to_port=to_port,
892+
cidr_ip=cidr_ip
893+
)
894+
#assert
895+
client.revoke_security_group_ingress.assert_called_with(
896+
GroupId=requested_security_group_id,
897+
IpPermissions=[
898+
{
899+
'IpProtocol': ip_protocol,
900+
'FromPort': from_port,
901+
'ToPort': to_port,
902+
'IpRanges': [{'CidrIp': cidr_ip}],
903+
'UserIdGroupPairs': [{}]
904+
}
905+
]
906+
)

0 commit comments

Comments
 (0)