#! /usr/bin/python3

from pyVim.connect import SmartConnectNoSSL
from pyVmomi import vim
from pyVim import task as Task
import argparse



def retrieve_portgroup(content, pg_name):
    target_pg = None
    pg_view = content.viewManager.CreateContainerView(content.rootFolder,
        [vim.DistributedVirtualPortgroup],
        True)
    pgs = [pg for pg in pg_view.view]
    for pg in pgs:
        if pg.name == pg_name:
            target_pg = pg
            break
    if not target_pg:
        raise Exception("can not find the dvportgroup: %s" % (pg_name))
    return target_pg


def retrieve_dvs(content, dvs_name):
    target_dvs = None
    dvs_view = content.viewManager.CreateContainerView(content.rootFolder,
        [vim.VmwareDistributedVirtualSwitch],
        True);
    dvses = [dvs for dvs in dvs_view.view]
    for _dvs in dvses:
        if _dvs.name == dvs_name:
            target_dvs = _dvs
            break
    if not target_dvs:
        raise Exception('Can not find dvs:%s' % (dvs_name))
    return target_dvs

def create_portgroup(dvs, pg_name):
    assert(dvs)
    spec = vim.dvs.DistributedVirtualPortgroup.ConfigSpec()
    spec.name = pg_name
    spec.defaultPortConfig = vim.dvs.VmwareDistributedVirtualSwitch.VmwarePortConfigPolicy()
    spec.defaultPortConfig.vlan = vim.dvs.VmwareDistributedVirtualSwitch.VlanIdSpec()
    spec.defaultPortConfig.vlan.inherited = False
    spec.defaultPortConfig.vlan.vlanId = 0
    spec.defaultPortConfig.networkResourcePoolKey = vim.StringPolicy()
    spec.defaultPortConfig.networkResourcePoolKey.inherited = False
    spec.defaultPortConfig.networkResourcePoolKey.value = '-1'
    spec.type = 'earlyBinding'
    spec.numPorts = 8
    spec.autoExpand = True
    spec.vmVnicNetworkResourcePoolKey = '-1'
    task = dvs.AddDVPortgroup_Task([spec])
    Task.WaitForTask(task = task)
    print("successfully added portgroup:%s on dvs:%s" %(pg_name, dvs.name))

def enable_vlan_override(content, pg):
    configspec = vim.dvs.DistributedVirtualPortgroup.ConfigSpec()
    configspec.configVersion = pg.config.configVersion
    configspec.name = pg.config.name
    configspec.numPorts = pg.config.numPorts
    configspec.defaultPortConfig = pg.config.defaultPortConfig
    configspec.type = pg.config.type
    configspec.policy = pg.config.policy
    configspec.policy.vlanOverrideAllowed = True
    configspec.autoExpand = pg.config.autoExpand
    configspec.vmVnicNetworkResourcePoolKey = pg.config.vmVnicNetworkResourcePoolKey
    configspec.description = pg.config.description
    task = pg.ReconfigureDVPortgroup_Task(configspec)
    Task.WaitForTask(task = task)
    print("successfully enabled vlan override on portgroup:%s" %(pg.config.name))

def retrieve_port(dvs, port_key):
    arg = vim.dvs.PortCriteria()
    arg.portKey = [str(port_key)]
    dvports = dvs.FetchDVPorts(arg)
    assert(len(dvports) == 1 and dvports[0].key == str(port_key))
    return dvports[0]

def reconfigure_port(dvs, port_key, vlan_id):
    dvport = retrieve_port(dvs, port_key)
    spec = vim.dvs.DistributedVirtualPort.ConfigSpec()
    spec.operation = 'edit'
    spec.key = str(port_key)
    spec.setting = vim.dvs.VmwareDistributedVirtualSwitch.VmwarePortConfigPolicy()
    spec.setting.blocked = vim.BoolPolicy()
    spec.setting.blocked.inherited = True
    spec.setting.blocked.value = False
    spec.setting.vlan = vim.dvs.VmwareDistributedVirtualSwitch.VlanIdSpec()
    spec.setting.vlan.inherited = False
    spec.setting.vlan.vlanId = int(vlan_id)
    spec.configVersion = dvport.config.configVersion
    task = dvs.ReconfigureDVPort_Task([spec])
    Task.WaitForTask(task = task)
    print("successfully modified port %s 's vlan to:%s on dvs:%s" %(port_key, vlan_id, dvs.name))

def parse_arguments():
    parser = argparse.ArgumentParser(description='Process script arguments...')
    parser.add_argument('--host', type = str, required = True,  help = 'the IP or name of the vSphere Center')
    parser.add_argument('--port', type = int, default = 443)
    parser.add_argument('--user', type = str, help = 'username', default = 'Administrator@vsphere.local')
    parser.add_argument('--pwd', type = str, help = 'password', default = 'Admin!23')
    parser.add_argument('--dvs', type = str, required = True, help = 'the dvs name')
    parser.add_argument('--portgroup', type = str, help = 'the portgroup name')
    parser.add_argument('--key', type = int, help = 'port key')
    parser.add_argument('--vlan', type = int, help = 'the vlan id to be configured')
    parser.add_argument('--obj', type = str, required = True, help = 'the operation in[portgroup, port]')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_arguments()
    instance = SmartConnectNoSSL(host = args.host,
        user = args.user,
        pwd = args.pwd,
        port = args.port)
    content = instance.RetrieveContent()
    dvs = retrieve_dvs(content, args.dvs)
    if args.obj == 'portgroup':
        create_portgroup(dvs, args.portgroup)
        pg = retrieve_portgroup(content, args.portgroup)
        enable_vlan_override(content, pg)
    elif args.obj == 'port':
        reconfigure_port(dvs, args.key, args.vlan)