From 107788993f5dbd01d340e79b567370c1774bf9c0 Mon Sep 17 00:00:00 2001
From: Philip Schaten <philip.schaten@stud.uni-goettingen.de>
Date: Thu, 14 May 2020 00:24:55 +0200
Subject: [PATCH] Add functionality to change Access Point Port VLANs

---
 pyunifi/controller.py | 85 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/pyunifi/controller.py b/pyunifi/controller.py
index 2334312..982e57b 100644
--- a/pyunifi/controller.py
+++ b/pyunifi/controller.py
@@ -270,6 +270,91 @@ def get_switch_port_overrides(self, target_mac):
         log.debug('get_switch_port_overrides(%s)', target_mac)
         return self.get_device_stat(target_mac)['port_overrides']
 
+    def get_networks(self):
+        """Returns unfiltered list of all configured networks
+        :returns: networks known to the unifi controller
+        :rtype: list()
+        """
+        log.debug('get_networks()')
+        return self._api_read("rest/networkconf")
+
+    def _get_portconf_id_list(self):
+        """VLANs on an(y) AP port are mapped to a "portconf_id".
+        This helper method gets the mapping vlan id -> portconf_id
+        :returns: dictionary vlan_id -> portconf_id
+        :rtype: dict()
+        """
+        nets = self.get_networks()
+        portconf_id_list = {}
+        for net in nets:
+            if 'vlan' in net:
+                portconf_id_list[int(net['vlan'])] = net['_id']
+        return portconf_id_list
+
+    def _ap_port_vlan(self, target_mac, port_idx, portconf_id):
+        """Helper method to set the VLAN on the user ports of UAP-IW Devices.
+        Generates the json parameters including all unchanged settings.
+
+        :param target_mac: MAC address of the AP.
+        :type target_mac: str()
+        :param port_idx: Port ID to target
+        :type port_idx: int()
+        :param portconf_id: portconf_id to set.
+        :type portconf_id: str()
+        :returns: { 'port_overrides': [ { 'port_idx': int(),
+            'portconf': str, 'name': str } ] }
+        :rtype: dict( list( dict() ) )
+        """
+        log.debug('_ap_port_vlan(%s, %s, %s)', target_mac, port_idx, portconf_id)
+        device_stat = self.get_device_stat(target_mac)
+        device_id = device_stat['_id']
+        overrides = device_stat['port_overrides']
+        found = False
+        for i in range(0, len(overrides)):
+            if overrides[i]['port_idx'] == port_idx:
+                # Override already exists, update..
+                overrides[i]['portconf_id'] = portconf_id
+                found = True
+                break
+        if not found:
+            log.error("No overrides for port %s on device %s. Are port vlans enabled/is this the right port?" % (str(port_idx), target_mac))
+            raise APIError(
+                "No overrides for port %s on device %s. Are port vlans enabled/is this the right port?" % (str(port_idx), target_mac)
+            )
+        return {"port_overrides": overrides, "device_id": device_id}
+
+    def ap_port_vlan(self, target_mac, port_idx, vlan, portconf_id_list=None):
+        """Changes the VLAN on the given UAP Port
+
+        :param target_mac: MAC address of the AP.
+        :type target_mac: str()
+        :param port_idx: Port ID to change
+        :type port_idx: int()
+        :param vlan: VLAN ID to set
+        :type vlan: int()
+        :param portconf_id_list: Optional mapping vlan -> portconf_id as generated by get_portconf_id_list.
+        Otherwise the controller is queried for every port change
+        :type portconf_id_list: dict()
+
+        :returns: API Response which is either the full device info or empty if nothing was changed
+        :rtype: list()
+        """
+        log.debug('ap_port_vlan(%s, %s, %s)', target_mac, port_idx, vlan)
+
+        if portconf_id_list is None:
+            portconf_id_list = self._get_portconf_id_list()
+        if vlan not in portconf_id_list:
+            log.error("VLAN %s unknown (no corresponding portconf_id was found)." % str(vlan))
+            raise APIError(
+                "VLAN %s unknown (no corresponding portconf_id was found)." % str(vlan)
+            )
+        # The portconf_id is the _id of the network w/ that VLAN, but + 1. !?
+        portconf_id = hex(int(portconf_id_list[vlan], 16)+1)[2:]
+        params = self._ap_port_vlan(target_mac, port_idx, portconf_id)
+        device_id = params['device_id']
+        del params['device_id']
+        return self._api_update('rest/device/' + device_id, params)
+
     def _switch_port_power(self, target_mac, port_idx, mode):
         """Helper method to set the given PoE mode the port/switch.