from ipaddress import ip_network, ip_address
from pulumi import ComponentResource, ResourceOptions, StackReference
import vdc

class HubProps:
    def __init__(
        self,
        azure_bastion: bool,
        forced_tunnel: bool,
        firewall_address_space: str,
        hub_address_space: str,
        location: str,
        peer: str,
        reference: StackReference,
        resource_group_name: str,
        separator: str,
        stack: str,
        subnets: [str, str, str],
        suffix: str,
        tags: [str, str],
    ):
        self.azure_bastion = azure_bastion
        self.forced_tunnel = forced_tunnel
        self.firewall_address_space = firewall_address_space
        self.hub_address_space = hub_address_space
        self.location = location
        self.peer = peer
        self.reference = reference
        self.resource_group_name = resource_group_name
        self.separator = separator
        self.stack = stack
        self.subnets = subnets
        self.suffix = suffix
        self.tags = tags

class Hub(ComponentResource):
    def __init__(self, name: str, props: HubProps, opts: ResourceOptions=None):
        super().__init__('vdc:network:Hub', name, {}, opts)

        # set required vdc variables before calling functions
        vdc.location = props.location
        vdc.resource_group_name = props.resource_group_name
        vdc.s = props.separator
        vdc.self = self
        vdc.suffix = props.suffix
        vdc.tags = props.tags

        # calculate the subnets in the firewall_address_space
        fwz_nw = ip_network(props.firewall_address_space)
        fwz_sn = fwz_nw.subnets(new_prefix=25) # two /26 subnets required
        fwx_nw = next(fwz_sn) # for Azure Firewall and Management subnets
        fwz_sn = fwz_nw.address_exclude(fwx_nw) # consolidate remainder
        dmz_nw = next(fwz_sn) # largest remaining subnet for DMZ
        fwx_sn = fwx_nw.subnets(new_prefix=26) # split the /25 into two /26
        fws_nw = next(fwx_sn) # AzureFirewallSubnet
        fwm_nw = next(fwx_sn) # AzureFirewallManagementSubnet

        # calculate the subnets in the hub_address_space
        hub_nw = ip_network(props.hub_address_space)
        if hub_nw.prefixlen < 20: # split evenly between subnets and hosts
            sub_diff = int((hub_nw.max_prefixlen - hub_nw.prefixlen) / 2)
        else:
            sub_diff = 25 - hub_nw.prefixlen # minimum /25 subnet
        subnets = hub_nw.subnets(prefixlen_diff=sub_diff)
        next_sn = next(subnets) # first subnet reserved for special uses
        first_sn = next_sn.subnets(new_prefix=26) # split it into /26 subnets
        gws_nw = next(first_sn) # GatewaySubnet /26
        rem_nw = next(first_sn) # at least one more /26 subnet, perhaps more
        rem_sn = rem_nw.subnets(new_prefix=27) # only need /27 save the rest
        abs_nw = next(rem_sn) # AzureBastionSubnet /27 or greater

        # cast repeatedly referenced networks to strings
        dmz_ar = str(dmz_nw)
        gws_ar = str(gws_nw)

        # set the separator to be used in resource names
        s = props.separator

        # Azure Virtual Network to which spokes will be peered
        # separate address spaces to simplify custom routing
        hub = vdc.virtual_network(name, [
                props.firewall_address_space,
                props.hub_address_space,
            ],
        )

        # AzureFirewallManagementSubnet and Route Table
        # https://docs.microsoft.com/en-us/azure/firewall/forced-tunneling
        hub_fwm_rt = vdc.route_table(
            stem = f'{name}{s}fwm',
            disable_bgp_route_propagation = True, #required
        )
        # only a default route to the Internet is permitted
        hub_fwm_dg = vdc.route_to_internet(
            stem = f'fwm{s}internet',
            route_table_name = hub_fwm_rt.name,
        )
        hub_fwm_sn = vdc.subnet_special(
            stem = f'{name}{s}fwm',
            name = 'AzureFirewallManagementSubnet', # name required
            virtual_network_name = hub.name,
            address_prefix = str(fwm_nw),
            route_table_id = hub_fwm_rt.id,
            depends_on = [hub, hub_fwm_rt, hub_fwm_dg],
        )

        # AzureFirewallSubnet and Route Table
        hub_fw_rt = vdc.route_table(
            stem = f'{name}{s}fw',
            disable_bgp_route_propagation = False,
        )
        # default route either direct to Internet or forced tunnel
        # turn off SNAT if the next_hop_ip_address is public
        # https://docs.microsoft.com/en-us/azure/firewall/snat-private-range
        private_ranges = 'IANAPrivateRanges'
        if not props.forced_tunnel:
            hub_fw_dg = vdc.route_to_internet(
                stem = f'fw{s}internet',
                route_table_name = hub_fw_rt.name,
            )
        else:
            hub_fw_dg = vdc.route_to_virtual_appliance(
                stem = f'fw{s}tunnel',
                route_table_name = hub_fw_rt.name,
                address_prefix = '0.0.0.0/0',
                next_hop_ip_address = props.forced_tunnel,
            )
            ft_ip = ip_address(props.forced_tunnel)
            if not ft_ip.is_private:
                private_ranges = '0.0.0.0/0'
        hub_fw_sn = vdc.subnet_special(
            stem = f'{name}{s}fw',
            name = 'AzureFirewallSubnet', # name required
            virtual_network_name = hub.name,
            address_prefix = str(fws_nw),
            route_table_id = hub_fw_rt.id,
            depends_on = [hub, hub_fw_rt, hub_fw_dg],
        )

        # Azure Firewall
        hub_fw = vdc.firewall(
            stem = name,
            fw_sn_id = hub_fw_sn.id,
            fwm_sn_id = hub_fwm_sn.id,
            private_ranges = private_ranges,
            depends_on = [hub_fw_sn, hub_fwm_sn],
        )

        # wait for the private ip address of the firewall to become available
        hub_fw_ip = hub_fw.ip_configurations.apply(
            lambda ipc: ipc[0].private_ip_address
        )
        # It is very important to ensure that there is never a route with an
        # address_prefix which covers the AzureFirewallSubnet.

        # DMZ subnet and Route Table
        hub_dmz_rt = vdc.route_table(
            stem = f'{name}{s}dmz',
            disable_bgp_route_propagation = True,
            depends_on = [hub_fw],
        )
        # default route from DMZ via the firewall
        hub_dmz_dg = vdc.route_to_virtual_appliance(
            stem = f'dmz{s}dg',
            route_table_name = hub_dmz_rt.name,
            address_prefix = '0.0.0.0/0',
            next_hop_ip_address = hub_fw_ip,
        )
        # redirect intra-DMZ traffic via the firewall
        hub_dmz_dmz = vdc.route_to_virtual_appliance(
            stem = f'dmz{s}dmz',
            route_table_name = hub_dmz_rt.name,
            address_prefix = dmz_ar,
            next_hop_ip_address = hub_fw_ip,
        )
        # redirect traffic from DMZ to hub via the firewall
        hub_dmz_hub = vdc.route_to_virtual_appliance(
            stem = f'dmz{s}hub',
            route_table_name = hub_dmz_rt.name,
            address_prefix = props.hub_address_space,
            next_hop_ip_address = hub_fw_ip,
        )
        hub_dmz_sn = vdc.subnet_special( #ToDo add NSG
            stem = f'{name}{s}dmz',
            name = 'DMZ', # name not required but preferred
            virtual_network_name = hub.name,
            address_prefix = dmz_ar,
            route_table_id = hub_dmz_rt.id,
            depends_on = [hub_dmz_rt, hub_dmz_dg, hub_dmz_dmz, hub_dmz_hub],
        )

        # GatewaySubnet and Route Table
        hub_gw_rt = vdc.route_table(
            stem = f'{name}{s}gw',
            disable_bgp_route_propagation = False,
            depends_on = [hub_dmz_sn],
        )
        # protect intra-GatewaySubnet traffic from being redirected:
        hub_gw_gw = vdc.route_to_virtual_network(
            stem = f'gw{s}gw',
            route_table_name = hub_gw_rt.name,
            address_prefix = gws_ar,
        )
        # redirect traffic from gateways to DMZ via firewall
        hub_gw_dmz = vdc.route_to_virtual_appliance(
            stem = f'gw{s}dmz',
            route_table_name = hub_gw_rt.name,
            address_prefix = dmz_ar,
            next_hop_ip_address = hub_fw_ip,
        )
        # redirect traffic from gateways to hub via firewall
        hub_gw_hub = vdc.route_to_virtual_appliance(
            stem = f'gw{s}hub',
            route_table_name = hub_gw_rt.name,
            address_prefix = props.hub_address_space,
            next_hop_ip_address = hub_fw_ip,
        )
        hub_gw_sn = vdc.subnet_special(
            stem = f'{name}{s}gw',
            name = 'GatewaySubnet', # name required
            virtual_network_name = hub.name,
            address_prefix = gws_ar,
            route_table_id = hub_gw_rt.id,
            depends_on = [hub_gw_rt, hub_gw_gw, hub_gw_dmz, hub_gw_hub],
        )

        # VPN Gateway
        hub_vpn_gw = vdc.vpn_gateway(
            stem = name,
            subnet_id = hub_gw_sn.id,
            depends_on = [hub_gw_sn],
        )

        # ExpressRoute Gateway
        hub_er_gw = vdc.expressroute_gateway(
            stem = name,
            subnet_id = hub_gw_sn.id,
            depends_on = [hub_gw_sn],
        )

        # Route Table to be associated with all hub shared services subnets
        hub_ss_rt = vdc.route_table(
            stem = f'{name}{s}ss',
            disable_bgp_route_propagation = True,
            depends_on = [hub_er_gw, hub_vpn_gw],
        )
        # default route from hub via the firewall
        hub_ss_dg = vdc.route_to_virtual_appliance(
            stem = f'ss{s}dg',
            route_table_name = hub_ss_rt.name,
            address_prefix = '0.0.0.0/0',
            next_hop_ip_address = hub_fw_ip,
        )
        # redirect traffic from hub to DMZ via the firewall
        hub_ss_dmz = vdc.route_to_virtual_appliance(
            stem = f'ss{s}dmz',
            route_table_name = hub_ss_rt.name,
            address_prefix = dmz_ar,
            next_hop_ip_address = hub_fw_ip,
        )
        # redirect traffic from hub to gateways via the firewall
        hub_ss_gw = vdc.route_to_virtual_appliance(
            stem = f'ss{s}gw',
            route_table_name = hub_ss_rt.name,
            address_prefix = gws_ar,
            next_hop_ip_address = hub_fw_ip,
        )
        # shared services subnets starting with the second subnet
        for subnet in props.subnets:
            next_sn = next(subnets)
            hub_sn = vdc.subnet( #ToDo add NSG
                stem = f'{name}{s}{subnet[0]}',
                virtual_network_name = hub.name,
                address_prefix = str(next_sn),
                route_table_id = hub_ss_rt.id,
                depends_on = [hub_ss_rt, hub_ss_dg, hub_ss_dmz, hub_ss_gw],
            )

        # Azure Bastion subnet and host (optional)
        if props.azure_bastion:
            hub_ab = vdc.bastion_host(
                stem = name,
                virtual_network_name = hub.name,
                address_prefix = str(abs_nw),
                depends_on = [hub_er_gw, hub_vpn_gw],
            )

        # VNet Peering between stacks using StackReference (optional)
        if props.peer:
            peer_hub_id = props.reference.get_output('hub_id')
            # VNet Peering (Global) in one direction from stack to peer
            hub_hub = vdc.vnet_peering(
                stem = props.stack,
                virtual_network_name = hub.name,
                peer = props.peer,
                remote_virtual_network_id = peer_hub_id,
                allow_forwarded_traffic = True,
                allow_gateway_transit = False, # as both hubs have gateways
            )
            # need to invalidate system routes created by VNet Peering
            peer_dmz_ar = props.reference.get_output('dmz_ar') 
            peer_fw_ip = props.reference.get_output('fw_ip')
            peer_hub_as = props.reference.get_output('hub_as')   
            for route in [
                (f'dmz{s}{props.peer}{s}dmz', hub_dmz_rt.name, peer_dmz_ar),
                (f'dmz{s}{props.peer}{s}hub', hub_dmz_rt.name, peer_hub_as),
                (f'gw{s}{props.peer}{s}dmz', hub_gw_rt.name, peer_dmz_ar),
                (f'gw{s}{props.peer}{s}hub', hub_gw_rt.name, peer_hub_as),
                (f'ss{s}{props.peer}{s}dmz', hub_ss_rt.name, peer_dmz_ar),
                (f'ss{s}{props.peer}{s}hub', hub_ss_rt.name, peer_hub_as),
            ]:
                vdc.route_to_virtual_appliance(
                    stem = route[0],
                    route_table_name = route[1],
                    address_prefix = route[2],
                    next_hop_ip_address = peer_fw_ip,
                )

        # assign properties to hub including from child resources
        self.address_space = props.hub_address_space # used for routes to the hub
        self.dmz_ar = dmz_ar # used for routes to the hub
        self.dmz_rt_name = hub_dmz_rt.name # used to add routes to spokes
        self.er_gw = hub_er_gw # needed prior to VNet Peering from spokes
        self.fw = hub_fw # needed prior to VNet Peering from spokes
        self.fw_ip = hub_fw_ip # used for routes to the hub
        self.fw_rt_name = hub_fw_rt.name # used for route to the peered spokes
        self.gw_rt_name = hub_gw_rt.name # used to add routes to spokes
        self.id = hub.id # exported and used for stack and spoke peering
        self.location = hub.location # informational
        self.name = hub.name # exported and used for spoke peering
        self.peer = props.peer # informational
        self.resource_group_name = props.resource_group_name # informational
        self.subnets = hub.subnets # informational
        self.stack = props.stack # informational
        self.stem = name # used for VNet Peering from spokes
        self.ss_rt_name = hub_ss_rt.name # used to add routes to spokes
        self.tags = props.tags # informational
        self.vpn_gw = hub_vpn_gw # needed prior to VNet Peering from spokes
        self.register_outputs({})
