chore(ec2): improve handling of ENIs (#3798)

This commit is contained in:
Sergio Garcia
2024-04-17 13:12:31 +02:00
committed by GitHub
parent a2f84a12ea
commit d79ec44e4c
3 changed files with 68 additions and 97 deletions
@@ -29,8 +29,7 @@ class EC2(AWSService):
self.__threading_call__(self.__describe_snapshots__)
self.__threading_call__(self.__determine_public_snapshots__, self.snapshots)
self.network_interfaces = []
self.__threading_call__(self.__describe_public_network_interfaces__)
self.__threading_call__(self.__describe_sg_network_interfaces__)
self.__threading_call__(self.__describe_network_interfaces__)
self.images = []
self.__threading_call__(self.__describe_images__)
self.volumes = []
@@ -243,7 +242,7 @@ class EC2(AWSService):
f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
def __describe_public_network_interfaces__(self, regional_client):
def __describe_network_interfaces__(self, regional_client):
try:
# Get Network Interfaces with Public IPs
describe_network_interfaces_paginator = regional_client.get_paginator(
@@ -251,47 +250,45 @@ class EC2(AWSService):
)
for page in describe_network_interfaces_paginator.paginate():
for interface in page["NetworkInterfaces"]:
if interface.get("Association"):
self.network_interfaces.append(
NetworkInterface(
public_ip=interface["Association"]["PublicIp"],
type=interface["InterfaceType"],
private_ip=interface["PrivateIpAddress"],
subnet_id=interface["SubnetId"],
vpc_id=interface["VpcId"],
region=regional_client.region,
tags=interface.get("TagSet"),
)
)
eni = NetworkInterface(
id=interface["NetworkInterfaceId"],
association=interface.get("Association", {}),
attachment=interface.get("Attachment", {}),
private_ip=interface["PrivateIpAddress"],
type=interface["InterfaceType"],
subnet_id=interface["SubnetId"],
vpc_id=interface["VpcId"],
region=regional_client.region,
tags=interface.get("TagSet"),
)
self.network_interfaces.append(eni)
# Add Network Interface to Security Group
# 'Groups': [
# {
# 'GroupId': 'sg-xxxxx',
# 'GroupName': 'default',
# },
# ],
self.__add_network_interfaces_to_security_groups__(
eni, interface.get("Groups", [])
)
except Exception as error:
logger.error(
f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
def __describe_sg_network_interfaces__(self, regional_client):
def __add_network_interfaces_to_security_groups__(
self, interface, interface_security_groups
):
try:
# Get Network Interfaces for Security Groups
for sg in self.security_groups:
regional_client = self.regional_clients[sg.region]
describe_network_interfaces_paginator = regional_client.get_paginator(
"describe_network_interfaces"
)
for page in describe_network_interfaces_paginator.paginate(
Filters=[
{
"Name": "group-id",
"Values": [
sg.id,
],
},
],
):
for interface in page["NetworkInterfaces"]:
sg.network_interfaces.append(interface["NetworkInterfaceId"])
for sg in interface_security_groups:
for security_group in self.security_groups:
if security_group.id == sg["GroupId"]:
security_group.network_interfaces.append(interface)
except Exception as error:
logger.error(
f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
def __get_instance_user_data__(self, instance):
@@ -449,6 +446,18 @@ class Volume(BaseModel):
tags: Optional[list] = []
class NetworkInterface(BaseModel):
id: str
association: dict
attachment: dict
private_ip: str
type: str
subnet_id: str
vpc_id: str
region: str
tags: Optional[list] = []
class SecurityGroup(BaseModel):
name: str
arn: str
@@ -457,7 +466,7 @@ class SecurityGroup(BaseModel):
vpc_id: str
public_ports: bool
associated_sgs: list
network_interfaces: list[str] = []
network_interfaces: list[NetworkInterface] = []
ingress_rules: list[dict]
egress_rules: list[dict]
tags: Optional[list] = []
@@ -472,16 +481,6 @@ class NetworkACL(BaseModel):
tags: Optional[list] = []
class NetworkInterface(BaseModel):
public_ip: str
private_ip: str
type: str
subnet_id: str
vpc_id: str
region: str
tags: Optional[list] = []
class ElasticIP(BaseModel):
public_ip: Optional[str]
association_id: Optional[str]
@@ -18,9 +18,13 @@ class route53_dangling_ip_subdomain_takeover(Check):
# Gather Elastic IPs and Network Interfaces Public IPs inside the AWS Account
public_ips = []
public_ips.extend([eip.public_ip for eip in ec2_client.elastic_ips])
public_ips.extend(
[interface.public_ip for interface in ec2_client.network_interfaces]
)
# Add public IPs from Network Interfaces
for network_interface in ec2_client.network_interfaces:
if (
network_interface.association
and network_interface.association.get("PublicIp")
):
public_ips.append(network_interface.association.get("PublicIp"))
for record in record_set.records:
# Check if record is an IP Address
if validate_ip_address(record):
@@ -369,51 +369,7 @@ class Test_EC2_Service:
# Test EC2 Describe Network Interfaces
@mock_aws
def test__describe_sg_network_interfaces__(self):
# Generate EC2 Client
ec2_client = client("ec2", region_name=AWS_REGION_US_EAST_1)
ec2_resource = resource("ec2", region_name=AWS_REGION_US_EAST_1)
# Create VPC, Subnet, SecurityGroup and Network Interface
vpc = ec2_resource.create_vpc(CidrBlock="10.0.0.0/16")
subnet = ec2_resource.create_subnet(VpcId=vpc.id, CidrBlock="10.0.0.0/18")
sg = ec2_resource.create_security_group(
GroupName="test-securitygroup", Description="n/a"
)
eni_id = subnet.create_network_interface(Groups=[sg.id]).id
ec2_client.modify_network_interface_attribute(
NetworkInterfaceId=eni_id, Groups=[sg.id]
)
# EC2 client for this test class
aws_provider = set_mocked_aws_provider(
[AWS_REGION_EU_WEST_1, AWS_REGION_US_EAST_1]
)
ec2 = EC2(aws_provider)
assert sg.id in str(ec2.security_groups)
for security_group in ec2.security_groups:
if security_group.id == sg.id:
assert security_group.name == "test-securitygroup"
assert (
security_group.arn
== f"arn:{aws_provider.identity.partition}:ec2:{AWS_REGION_US_EAST_1}:{AWS_ACCOUNT_NUMBER}:security-group/{security_group.id}"
)
assert re.match(r"sg-[0-9a-z]{17}", security_group.id)
assert security_group.region == AWS_REGION_US_EAST_1
assert eni_id in security_group.network_interfaces
assert security_group.ingress_rules == []
assert security_group.egress_rules == [
{
"IpProtocol": "-1",
"IpRanges": [{"CidrIp": "0.0.0.0/0"}],
"Ipv6Ranges": [],
"PrefixListIds": [],
"UserIdGroupPairs": [],
}
]
@mock_aws
def test__describe_public_network_interfaces__(self):
def test__describe_network_interfaces__(self):
# Generate EC2 Client
ec2_client = client("ec2", region_name=AWS_REGION_US_EAST_1)
ec2_resource = resource("ec2", region_name=AWS_REGION_US_EAST_1)
@@ -436,23 +392,35 @@ class Test_EC2_Service:
ec2_client.associate_address(
NetworkInterfaceId=eni.id, AllocationId=eip["AllocationId"]
)
# Attach ENI to Instance
ec2_resource.create_instances(
ImageId=EXAMPLE_AMI_ID,
MinCount=1,
MaxCount=1,
NetworkInterfaces=[{"DeviceIndex": 0, "NetworkInterfaceId": eni.id}],
)[0]
# EC2 client for this test class
aws_provider = set_mocked_aws_provider(
[AWS_REGION_EU_WEST_1, AWS_REGION_US_EAST_1]
)
ec2 = EC2(aws_provider)
assert len(ec2.network_interfaces) == 1
assert ec2.network_interfaces[0].public_ip == eip["PublicIp"]
assert ec2.network_interfaces[0].association
assert ec2.network_interfaces[0].attachment
assert ec2.network_interfaces[0].id == eni.id
assert ec2.network_interfaces[0].private_ip == eni.private_ip_address
assert ec2.network_interfaces[0].type == eni.interface_type
assert ec2.network_interfaces[0].subnet_id == subnet.id
assert ec2.network_interfaces[0].type == eni.interface_type
assert ec2.network_interfaces[0].vpc_id == vpc.id
assert ec2.network_interfaces[0].region == AWS_REGION_US_EAST_1
assert ec2.network_interfaces[0].tags == [
{"Key": "string", "Value": "string"},
]
# Check if ENI was added to security group
for sg in ec2.security_groups:
if sg.id == eni.groups[0]["GroupId"]:
assert sg.network_interfaces == ec2.network_interfaces
# Test EC2 Describe Images
@mock_aws