InboundSecurityGroupRules

from typing import List, Mapping, Optional, Sequence

import pulumi
import pulumi_aws as aws

from . import VPC_ID


class SecurityGroup(aws.ec2.SecurityGroup):
    """Creates a security group"""

    def __init__(
        self,
        name: str,
        description: str,
        vpc_id: str = VPC_ID,
        ingress: Optional[Sequence[aws.ec2.SecurityGroupIngressArgs]] = None,
        egress: Optional[Sequence[aws.ec2.SecurityGroupEgressArgs]] = None,
        tags: Optional[Mapping[str, str]] = None,
        opts: Optional[pulumi.ResourceOptions] = None,
    ) -> None:

        # Append "-sg" to name if not already present
        if name.endswith("-sg"):
            sg_name = name
        else:
            sg_name = f"{name}-sg"

        super().__init__(
            resource_name=sg_name,
            name=sg_name,
            description=description,
            vpc_id=vpc_id,
            ingress=ingress,
            egress=egress,
            opts=opts,
            tags={"Name": sg_name},
        )
        pulumi.export(name, {"security_group_id": self.id})


class CIDRSecurityGroupRule(aws.ec2.SecurityGroupRule):
    """Creates a security group rule to/from a CIDR"""

    def __init__(
        self,
        name: str,
        security_group_id: str,
        from_port: int,
        to_port: int,
        protocol: str,
        type: str,
        cidr_blocks: Optional[List[str]] = None,
        opts: Optional[pulumi.ResourceOptions] = None,
    ) -> None:
        sgr_name = f"{name}-sgr"
        super().__init__(
            resource_name=sgr_name,
            security_group_id=security_group_id,
            type=type,
            from_port=from_port,
            to_port=to_port,
            protocol=protocol,
            cidr_blocks=cidr_blocks,
        )
        pulumi.export(name, {"security_group_rule_id": self.id})


class SGSecurityGroupRule(aws.ec2.SecurityGroupRule):
    """Creates a security group rule to/from a security group"""

    def __init__(
        self,
        name: str,
        security_group_id: str,
        from_port: int,
        to_port: int,
        protocol: str,
        type: str,
        source_security_group_id: Optional[str] = None,
        opts: Optional[pulumi.ResourceOptions] = None,
    ) -> None:
        sgr_name = f"{name}-sgr"
        super().__init__(
            resource_name=sgr_name,
            security_group_id=security_group_id,
            type=type,
            from_port=from_port,
            to_port=to_port,
            protocol=protocol,
            source_security_group_id=source_security_group_id,
        )
        pulumi.export(name, {"security_group_rule_id": self.id})


class InboundSecurityGroupRules:
    """Creates security group rules for one TCP port from one or more CIDRs and/or one security group"""

    def __init__(
        self,
        parent_sg_id: str,
        port: int,
        source_cidrs: Optional[List[str]] = None,
        source_sg_id: Optional[str] = None,
        opts: Optional[pulumi.ResourceOptions] = None,
    ) -> None:

        if source_cidrs is not None:

            if isinstance(source_cidrs, str):
                subnets = [source_cidrs]
                rule_name_suffix = source_cidrs
            else:
                subnets = [subnet.value for subnet in source_cidrs]
                rule_name_suffix = source_cidrs.__name__

            self.cidr_rule = CIDRSecurityGroupRule(
                name=f"{parent_sg_id}-{rule_name_suffix}",
                cidr_blocks=subnets,
                security_group_id=parent_sg_id,
                from_port=port,
                to_port=port,
                opts=opts,
                type="ingress",
                protocol="tcp",
            )

        if source_sg_id is not None:
            self.sg_rule = SGSecurityGroupRule(
                name=f"{parent_sg_id}-{source_sg_id}",
                source_security_group_id=source_sg_id,
                security_group_id=parent_sg_id,
                from_port=port,
                to_port=port,
                opts=opts,
                type="ingress",
                protocol="tcp",
            )

        if source_cidrs is None and source_sg_id is None:
            raise ValueError("Must provide at least one of source_cidrs or source_sg_id")


class OutboundSecurityGroupRules:
    """Creates security group rules for one TCP port to one or more CIDRs and/or one security group"""

    def __init__(
        self,
        parent_sg_id: str,
        port: int,
        destination_cidrs: Optional[List[str]] = None,
        destination_sg_id: Optional[str] = None,
        opts: Optional[pulumi.ResourceOptions] = None,
    ) -> None:

        if destination_cidrs is not None:
            if isinstance(destination_cidrs, str):
                subnets = [destination_cidrs]
                rule_name_suffix = destination_cidrs
            else:
                subnets = [subnet.value for subnet in destination_cidrs]
                rule_name_suffix = destination_cidrs.__name__

            self.cidr_rule = CIDRSecurityGroupRule(
                name=f"{parent_sg_id}-{rule_name_suffix}",
                cidr_blocks=subnets,
                security_group_id=parent_sg_id,
                from_port=port,
                to_port=port,
                type="egress",
                protocol="tcp",
            )

        if destination_sg_id is not None:
            self.sg_rule = SGSecurityGroupRule(
                name=f"{parent_sg_id}-{destination_sg_id}",
                source_security_group_id=destination_sg_id,
                security_group_id=parent_sg_id,
                from_port=port,
                to_port=port,
                type="egress",
                protocol="tcp",
            )

        if destination_cidrs is None and destination_sg_id is None:
            raise ValueError("Must provide at least one of destination_cidrs or destination_sg_id")

InboundSecurityGroupRule