refactor

1 files changed, 129 insertions, 113 deletions

diff --git a/tools/netbox/scripts/create-switch/create-switch-tg25.py b/tools/netbox/scripts/create-switch/create-switch-tg25.py
index 540ae84..5207c18 100644
--- a/tools/netbox/scripts/create-switch/create-switch-tg25.py
+++ b/tools/netbox/scripts/create-switch/create-switch-tg25.py
@@ -42,11 +42,13 @@ UPLINK_TYPES = (
     (InterfaceTypeChoices.TYPE_10GE_FIXED, '10G RJ45')
 )
 
+
 def generatePrefix(prefix, length):
     firstPrefix = prefix.get_first_available_prefix()
     out = list(firstPrefix.subnet(length, count=1))[0]
     return out
 
+
 class CreateSwitch(Script):
     class Meta:
         name = "Create Switch"
@@ -57,46 +59,46 @@ class CreateSwitch(Script):
         scheduling_enabled = False
 
     switch_name = StringVar(
-        description = "Switch name. Remember, e = access switch, d = distro switch",
-        required = True,
-        default = "e1.test" # default during development
+        description="Switch name. Remember, e = access switch, d = distro switch",
+        required=True,
+        default="e1.test"  # default during development
     )
     device_type = ObjectVar(
-        description = "Device model",
-        model = DeviceType,
-        default = DEFAULT_DEVICE_TYPE.id,
+        description="Device model",
+        model=DeviceType,
+        default=DEFAULT_DEVICE_TYPE.id,
     )
     device_role = ObjectVar(
-        description = "Device role",
-        model = DeviceRole,
-        default = DEFAULT_DEVICE_ROLE.id,
+        description="Device role",
+        model=DeviceRole,
+        default=DEFAULT_DEVICE_ROLE.id,
     )
     location = ObjectVar(
-        model = Location,
-        required = True,
-        default = Location.objects.get(name="Ringen") # Default during development
+        model=Location,
+        required=True,
+        default=Location.objects.get(name="Ringen")  # Default during development
     )
     uplink_type = ChoiceVar(
-        label = 'Uplink Type',
-        required = True,
-        description = "What type of interface should this switch be delivered on",
-        choices = UPLINK_TYPES,
-        default = InterfaceTypeChoices.TYPE_1GE_FIXED
+        label='Uplink Type',
+        required=True,
+        description="What type of interface should this switch be delivered on",
+        choices=UPLINK_TYPES,
+        default=InterfaceTypeChoices.TYPE_1GE_FIXED
     )
     destination_device_a = ObjectVar(
-        description = "Uplink device (A)",
-        required = True,
-        model = Device,
-        query_params = {
+        description="Uplink device (A)",
+        required=True,
+        model=Device,
+        query_params={
             'role': [DEVICE_ROLE_LEAF, DEVICE_ROLE_DISTRO, DEVICE_ROLE_UTSKUTT_DISTRO],
         },
-        default = DEFAULT_UPLINK_SWITCH
+        default=DEFAULT_UPLINK_SWITCH
     )
     destination_device_b = ObjectVar(
-        description = "If connected to leaf pair - Uplink device (B)",
-        required = False,
-        model = Device,
-        query_params = {
+        description="If connected to leaf pair - Uplink device (B)",
+        required=False,
+        model=Device,
+        query_params={
             'role': [DEVICE_ROLE_LEAF],
         },
     )
@@ -115,131 +117,86 @@ class CreateSwitch(Script):
         if not data['switch_name'].startswith("e") and not data['switch_name'].startswith("d"):
             raise ValidationError("Switch name must start whit e or d")
 
-        switch = Device(
-            name = data['switch_name'],
-            device_type = data['device_type'],
-            location = data['location'],
-            role = data['device_role'],
-            site = DEFAULT_SITE
-        )
-        switch.save()
-        self.log_info("Created switch")
-
-        mgmt_interface_name = "irb"
-        if switch.device_type.model == "EX2200-48T-4G":
-            mgmt_interface_name = "vlan"
-
-        mgmt_vlan_interface = Interface.objects.create(
-            device=switch,
-            name=f"{mgmt_interface_name}.{FABRIC_V4_JUNIPER_MGMT_PREFIX.vlan.id}",
-            description = f'X: Mgmt',
-            type=InterfaceTypeChoices.TYPE_VIRTUAL,
-            mode=InterfaceModeChoices.MODE_TAGGED,
-        )
-
-        v4_mgmt_addr = IPAddress.objects.create(
-            address=FABRIC_V4_JUNIPER_MGMT_PREFIX.get_first_available_ip(),
-            dns_name=f"{switch.name}.{DEFAULT_TG_DNS_SUFFIX}"
+        switch = self.create_switch(data)
 
-        )
-        v6_mgmt_addr = IPAddress.objects.create(
-            address=FABRIC_V6_JUNIPER_MGMT_PREFIX.get_first_available_ip(),
-            dns_name=f"{switch.name}.{DEFAULT_TG_DNS_SUFFIX}"
-        )
-
-        mgmt_vlan_interface.ip_addresses.add(v4_mgmt_addr)
-        mgmt_vlan_interface.ip_addresses.add(v6_mgmt_addr)
-        switch.primary_ip4 = v4_mgmt_addr
-        switch.primary_ip6 = v6_mgmt_addr
-        switch.save()
-        self.log_info("Allocated and assigned mgmt addresses on switch")
+        vlan = self.create_vlan(switch)
 
-        vid = FABRIC_VLAN_GROUP.get_next_available_vid()
-        vlan = VLAN.objects.create(
-            name = switch.name,
-            group = FABRIC_VLAN_GROUP,
-            role = FABRIC_CLIENTS_ROLE,
-            vid = vid
-        )
-        vlan.save()
-        self.log_info("Created VLAN")
+        v4_prefix, v6_prefix = self.allocate_prefixes(vlan)
 
-        interfaces = list(Interface.objects.filter(device=switch, type=InterfaceTypeChoices.TYPE_1GE_FIXED))
-        if len(interfaces) == 0:
-            self.log_error(f"no interfaces found")
+        self.set_traffic_vlan(switch, vlan)
 
-        for interface in interfaces:
-            if interface.name in UPLINK_PORTS.get(switch.device_type.model, []):
-                continue
-            interface.mode = 'access'
-            interface.untagged_vlan = vlan
-            interface.description = "C: Clients"
-            interface.save()
+        self.connect_switch(data, switch, vlan)
 
-        self.log_info("Configured traffic vlan on all client ports")
+        self.log_success(f"✅ Script completed successfully.")
+        self.log_success(f"🔗 Switch:     <a href=\"{switch.get_absolute_url()}\">{switch}</a>")
+        self.log_success(f"🔗 v6 Prefix:  <a href=\"{v6_prefix.get_absolute_url()}\">{v6_prefix}</a>")
+        self.log_success(f"🔗 v4 Prefix:  <a href=\"{v4_prefix.get_absolute_url()}\">{v4_prefix}</a>")
+        self.log_success(f"🔗 VLAN:       <a href=\"{vlan.get_absolute_url()}\">{vlan}</a>")
+        self.log_success(f"⚠️ <strong>️Fabric config must be deployed before switch can be fapped.</strong>")
 
+    def allocate_prefixes(self, vlan):
         v6_prefix = Prefix.objects.create(
-            prefix = generatePrefix(FABRIC_V6_CLIENTS_PREFIX, 64),
-            status = PrefixStatusChoices.STATUS_ACTIVE,
-            role = FABRIC_CLIENTS_ROLE,
-            vlan = vlan
+            prefix=generatePrefix(FABRIC_V6_CLIENTS_PREFIX, 64),
+            status=PrefixStatusChoices.STATUS_ACTIVE,
+            role=FABRIC_CLIENTS_ROLE,
+            vlan=vlan
         )
-
         v4_prefix = Prefix.objects.create(
-            prefix = generatePrefix(FABRIC_V4_CLIENTS_PREFIX, 26),
-            status = PrefixStatusChoices.STATUS_ACTIVE,
-            role = FABRIC_CLIENTS_ROLE,
-            vlan = vlan
+            prefix=generatePrefix(FABRIC_V4_CLIENTS_PREFIX, 26),
+            status=PrefixStatusChoices.STATUS_ACTIVE,
+            role=FABRIC_CLIENTS_ROLE,
+            vlan=vlan
         )
-        self.log_info("Allocated traffic prefixes")
+        self.log_info("Created network. Created new VLAN and assigned prefixes")
+        return v4_prefix, v6_prefix
 
+    def connect_switch(self, data, switch, vlan):
         uplink_description = f"B: {data['destination_device_a'].name}"
         if data['destination_device_b']:
             uplink_description = f"B: {data['destination_device_a'].name} / {data['destination_device_b'].name} - ae{vlan.id}"
-
         uplink_ae = Interface.objects.create(
             device=switch,
             name="ae0",
-            description = uplink_description,
-            type = InterfaceTypeChoices.TYPE_LAG,
-            mode = InterfaceModeChoices.MODE_TAGGED,
+            description=uplink_description,
+            type=InterfaceTypeChoices.TYPE_LAG,
+            mode=InterfaceModeChoices.MODE_TAGGED,
         )
-
-
         if data['destination_device_a'].role == DEVICE_ROLE_UTSKUTT_DISTRO:
-            self.log_debug(f"{ data['destination_device_a']} is utskutt-distro")
+            self.log_debug(f"{data['destination_device_a']} is utskutt-distro")
             # TODO make sure we add traffic vlan on AE between distro and utskutt-distro as well.
-
         uplink_ae.tagged_vlans.add(FABRIC_V4_JUNIPER_MGMT_PREFIX.vlan.id)
         uplink_ae.tagged_vlans.add(vlan.id)
-
         ## We only need this if not connected to leaf (since they are provisioned using AVD)
         if data['destination_device_a'].role != DEVICE_ROLE_LEAF:
             destination_ae = Interface.objects.create(
                 device=data['destination_device_a'],
                 name=f"ae{vlan.id}",
-                description = f'B: {switch.name} ae0',
+                description=f'B: {switch.name} ae0',
                 type=InterfaceTypeChoices.TYPE_LAG,
                 mode=InterfaceModeChoices.MODE_TAGGED,
             )
             destination_ae.save()
             destination_ae.tagged_vlans.add(FABRIC_V4_JUNIPER_MGMT_PREFIX.vlan.id)
             destination_ae.tagged_vlans.add(vlan.id)
-            self.log_debug(f"Created destination AE <a href=\"{destination_ae.get_absolute_url()}\">{destination_ae}</a>")
+            self.log_debug(
+                f"Created destination AE <a href=\"{destination_ae.get_absolute_url()}\">{destination_ae}</a>")
 
             ## TODO support leaf pair
             num_uplinks = len(data['destination_interfaces'])
             uplink_interfaces = list(Interface.objects.filter(device=switch, type=data['uplink_type']))
             if len(uplink_interfaces) < 1:
-                raise AbortScript(f"You chose a device type without any {data['uplink_type']} interfaces! Pick another model :)")
+                raise AbortScript(
+                    f"You chose a device type without any {data['uplink_type']} interfaces! Pick another model :)")
 
             interface_type = ContentType.objects.get_for_model(Interface)
-            interfaces_filtered = [interface for interface in uplink_interfaces if interface.name in UPLINK_PORTS.get(switch.device_type.model, [])]
+            interfaces_filtered = [interface for interface in uplink_interfaces if
+                                   interface.name in UPLINK_PORTS.get(switch.device_type.model, [])]
             for uplink_num in range(0, num_uplinks):
                 a_interface = data['destination_interfaces'][uplink_num]
                 b_interface = interfaces_filtered[uplink_num]
 
-                self.log_debug(f"Connecting: {data['destination_device_a']} - {a_interface} to {switch} - {b_interface}")
+                self.log_debug(
+                    f"Connecting: {data['destination_device_a']} - {a_interface} to {switch} - {b_interface}")
 
                 a_interface.description = f'G: {switch.name} {b_interface.name} (ae0)'
                 b_interface.description = f"G: {data['destination_device_a'].name} {a_interface.name} (ae{vlan.id})"
@@ -268,12 +225,71 @@ class CreateSwitch(Script):
                 cable._terminations_modified = True
                 cable.save()
 
+    def set_traffic_vlan(self, switch, vlan):
+        interfaces = list(Interface.objects.filter(device=switch, type=InterfaceTypeChoices.TYPE_1GE_FIXED))
+        if len(interfaces) == 0:
+            self.log_error(f"no interfaces found")
+
+        for interface in interfaces:
+            if interface.name in UPLINK_PORTS.get(switch.device_type.model, []):
+                continue
+            interface.mode = 'access'
+            interface.untagged_vlan = vlan
+            interface.description = "C: Clients"
+            interface.save()
+        self.log_info("Configured traffic vlan on all client ports")
+
+    def create_vlan(self, switch):
+        vid = FABRIC_VLAN_GROUP.get_next_available_vid()
+        vlan = VLAN.objects.create(
+            name=switch.name,
+            group=FABRIC_VLAN_GROUP,
+            role=FABRIC_CLIENTS_ROLE,
+            vid=vid
+        )
+        vlan.save()
+        self.log_info("Created VLAN")
+        return vlan
+
+    def create_switch(self, data):
+        switch = Device(
+            name=data['switch_name'],
+            device_type=data['device_type'],
+            location=data['location'],
+            role=data['device_role'],
+            site=DEFAULT_SITE
+        )
+        switch.save()
+
+        mgmt_interface_name = "irb"
+        if switch.device_type.model == "EX2200-48T-4G":
+            mgmt_interface_name = "vlan"
+
+        mgmt_vlan_interface = Interface.objects.create(
+            device=switch,
+            name=f"{mgmt_interface_name}.{FABRIC_V4_JUNIPER_MGMT_PREFIX.vlan.id}",
+            description=f'X: Mgmt',
+            type=InterfaceTypeChoices.TYPE_VIRTUAL,
+            mode=InterfaceModeChoices.MODE_TAGGED,
+        )
+        v4_mgmt_addr = IPAddress.objects.create(
+            address=FABRIC_V4_JUNIPER_MGMT_PREFIX.get_first_available_ip(),
+            dns_name=f"{switch.name}.{DEFAULT_TG_DNS_SUFFIX}"
+
+        )
+        v6_mgmt_addr = IPAddress.objects.create(
+            address=FABRIC_V6_JUNIPER_MGMT_PREFIX.get_first_available_ip(),
+            dns_name=f"{switch.name}.{DEFAULT_TG_DNS_SUFFIX}"
+        )
+        mgmt_vlan_interface.ip_addresses.add(v4_mgmt_addr)
+        mgmt_vlan_interface.ip_addresses.add(v6_mgmt_addr)
+        switch.primary_ip4 = v4_mgmt_addr
+        switch.primary_ip6 = v6_mgmt_addr
+        switch.save()
+
+        self.log_info("Created switch")
+        self.log_info("Allocated and assigned mgmt addresses on switch")
+        return switch
 
-        self.log_success(f"✅ Script completed successfully.")
-        self.log_success(f"🔗 Switch:     <a href=\"{switch.get_absolute_url()}\">{switch}</a>")
-        self.log_success(f"🔗 v6 Prefix:  <a href=\"{v6_prefix.get_absolute_url()}\">{v6_prefix}</a>")
-        self.log_success(f"🔗 v4 Prefix:  <a href=\"{v4_prefix.get_absolute_url()}\">{v4_prefix}</a>")
-        self.log_success(f"🔗 VLAN:       <a href=\"{vlan.get_absolute_url()}\">{vlan}</a>")
-        self.log_success(f"⚠️ <strong>️Fabric config must be deployed before switch can be fapped.</strong>")
 
 script = CreateSwitch