Update / Modify a Traffic Group of a Virtual Address using F5-SDK

I have enhanced the above code that will allow you to change a traffic group to any or all Virtual Addresses on a LTM unit.

from f5.bigip import ManagementRoot
import sys
from random import randrange
from pprint import pprint


 This module will read all the Virtual Servers from the bigip unit and create a dictionary e.g;
 VIP-name, IP address and Traffic-Group assigned

 VIPs_dict = { 'vs-client1.company.com-http':'['192.168.12.100','Traffic-group-1'],
               'vs-consumer.company.com-https':'['192.168.12.200','Traffic-group-1'],
               ....
            }
 Contributor : mshoaib@paciolan.com

 Environment
 Python : 3.6
 F5 SDK : 3.0.20
 LTM : 12.1.2 HF2

 Requirements :
 1. Create a folder name "data" in the same directory where this script will run. It will be used to save the current config
 2. Update the Variables section below according to your environment

def save_traffic_group_config(traffic_group_config):
  '''
   This function will save the current traffic groups config to a file
  '''
   This will create a file name "traffic_group_config_"
  new_file = open("data/traffic_group_conifg_"+str(randrange(1000)),'w')
  new_file.write("\n{}\t{}\t{}".format('Virtual','IP','Traffic Group'))
  for key in traffic_group_config.keys():
    new_file.write("\n{}\t{}\t{}".format(key,traffic_group_config[key][0],traffic_group_config[key][1]))
  new_file.close()
  print ("\nCurrent Virtual Addresses Traffic Group config have been saved ....")
 ------

def modify_traffic_group(virtual_server_name,virtual_address,bigip_mgmt):
  '''
   This function will take a name of virtual server and switch its traffic group from 1 to 2 OR 2 to 1
   Assumption : You have two traffic group named a) traffic-group-1 b) traffic-group-2
  '''

   Load a Virtual Address
  virt_address_info = bigip_mgmt.tm.ltm.virtual_address_s.virtual_address.load(name=virtual_address, partition='Common')

  print ("{} traffic group is {}".format(virtual_server_name,virt_address_info.trafficGroup))
  response = input("Type Y or yes to Change or  to Skip ... ")

   User Input to proceed
  if response.lower() == 'yes' or response.lower() == 'y':
    print (" Traffic Group Changed ")
  else:
    print (" No Change ")
    return
  if virt_address_info.trafficGroup == "/Common/traffic-group-1":
    new_TG = "/Common/traffic-group-2"
  else:
    new_TG = "/Common/traffic-group-1"

   Update the Traffic Group for that particular virtual server name / virtual address
  virt_address_info.modify(trafficGroup=new_TG)


 Variable Section
BigIP = "BigIP"
BigIP_username = "admin"
BigIP_password = "admin"


cur_virtuals_info = {}  Empty dict to collect Virtual Server Name, IP address & Traffic Group
tmp_virtual_address = []  it will store the IP of Virtual Addresses and help to avoid duplicates

 Connect to BigIP
f5_mgmt = ManagementRoot(BigIP, BigIP_username, BigIP_password)

 Gather Virtual Address and Virtual Server info
virt_addrs = f5_mgmt.tm.ltm.virtual_address_s.get_collection()
virt_srvs = f5_mgmt.tm.ltm.virtuals.get_collection()

 Fill up the Current Virtuals info
for v_s in virt_srvs:

  for v_a in virt_addrs:
     Compare Virtual Server IP and Virtual Address ( also an IP ) and if same then grab the traffic group and feed into the dict
    if v_a.name == v_s.destination.split('/')[2].split(':')[0]:  Match the IP address
      if len (cur_virtuals_info) == 0:
         Enter the first value in as is
        cur_virtuals_info[v_s.name] = [v_s.destination.split('/')[2].split(':')[0],v_a.trafficGroup.split('/')[2]]
         save the virtual address in tmp list to compare and avoid duplicates
        tmp_virtual_address.append(v_s.destination.split('/')[2].split(':')[0])
        continue
      else:
         Check if virtual address already exisit, if yes then skip
          if v_s.destination.split('/')[2].split(':')[0] in tmp_virtual_address:
             Virtual address already inserted, skip insertion
            continue
          else:
             Insert the Virtual Server / Address into the dict and in the temp list for comparison
            cur_virtuals_info[v_s.name] = [v_s.destination.split('/')[2].split(':')[0],v_a.trafficGroup.split('/')[2]]
            tmp_virtual_address.append(v_s.destination.split('/')[2].split(':')[0])
            continue

 Save the current Config
save_traffic_group_config(cur_virtuals_info)

for k in cur_virtuals_info.keys():
  modify_traffic_group(k,cur_virtuals_info[k][0],f5_mgmt)