Added VPN support on EX routers (#56)

* Added VPN support on EX routers * Adjusted E501 rules * Adjusted E501 rules
AlexandrErohin · Nov 19, 2024 · 0520699 · 0520699
1 parent 0beaeb7
commit 0520699
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 2 deletions.
diff --git a/test/test_client_ex.py b/test/test_client_ex.py
@@ -11,6 +11,8 @@
     IPv4DHCPLease,
     IPv4Status,
     ClientException,
+    VPNStatus,
+    VPN,
 )
 
 
@@ -334,6 +336,65 @@ def _request(self, url, method='POST', data_str=None, encrypt=False):
         self.assertEqual(check_data, '{"data":{"stack":"1,0,0,0,0,0","pstack":"0,0,0,0,0,0",'
                                      '"primaryEnable":"1"},"operation":"so","oid":"DEV2_ADT_WIFI_COMMON"}')
 
+    def test_get_vpn_status(self) -> None:
+        DEV2_OPENVPN = ('{"data":{"enable":"1","stack":"0,0,0,0,0,0"},'
+                        '"operation":"go","oid":"DEV2_OPENVPN","success":true}')
+        DEV2_PPTPVPN = ('{"data":{"enable":"0","stack":"0,0,0,0,0,0"},'
+                        '"operation":"go","oid":"DEV2_PPTPVPN","success":true}')
+        DEV2_OVPN_CLIENT = ('{"data":[{"connAct":"1","stack":"1,0,0,0,0,0"}, {"connAct":"1","stack":"2,0,0,0,0,0"},'
+                            '{"connAct":"0","stack":"3,0,0,0,0,0"}, {"connAct":"0","stack":"4,0,0,0,0,0"},'
+                            '{"connAct":"0","stack":"5,0,0,0,0,0"}, {"connAct":"0","stack":"6,0,0,0,0,0"},'
+                            '{"connAct":"0","stack":"7,0,0,0,0,0"}, {"connAct":"0","stack":"8,0,0,0,0,0"},'
+                            '{"connAct":"0","stack":"9,0,0,0,0,0"}, {"connAct":"0","stack":"10,0,0,0,0,0"}],'
+                            '"operation":"gl","oid":"DEV2_OVPN_CLIENT","success":true}')
+        DEV2_PVPN_CLIENT = ('{"data":[{"connAct":"0","stack":"1,0,0,0,0,0"}, {"connAct":"0","stack":"2,0,0,0,0,0"},'
+                            '{"connAct":"1","stack":"3,0,0,0,0,0"}, {"connAct":"0","stack":"4,0,0,0,0,0"},'
+                            '{"connAct":"0","stack":"5,0,0,0,0,0"}, {"connAct":"0","stack":"6,0,0,0,0,0"},'
+                            '{"connAct":"0","stack":"7,0,0,0,0,0"}, {"connAct":"0","stack":"8,0,0,0,0,0"},'
+                            '{"connAct":"0","stack":"9,0,0,0,0,0"}, {"connAct":"0","stack":"10,0,0,0,0,0"}],'
+                            '"operation":"gl","oid":"DEV2_PVPN_CLIENT","success":true}')
+
+        class TPLinkEXClientTest(TPLinkEXClient):
+            def _request(self, url, method='POST', data_str=None, encrypt=False):
+                if 'DEV2_OPENVPN' in data_str:
+                    return 200, DEV2_OPENVPN
+                elif 'DEV2_PPTPVPN' in data_str:
+                    return 200, DEV2_PPTPVPN
+                elif 'DEV2_OVPN_CLIENT' in data_str:
+                    return 200, DEV2_OVPN_CLIENT
+                elif 'DEV2_PVPN_CLIENT' in data_str:
+                    return 200, DEV2_PVPN_CLIENT
+                raise ClientException()
+
+        client = TPLinkEXClientTest('', '')
+        status = client.get_vpn_status()
+
+        self.assertIsInstance(status, VPNStatus)
+        self.assertEqual(status.openvpn_enable, True)
+        self.assertEqual(status.pptpvpn_enable, False)
+        self.assertEqual(status.openvpn_clients_total, 2)
+        self.assertEqual(status.pptpvpn_clients_total, 1)
+
+    def test_set_vpn(self) -> None:
+        response = '{"success":true, "errorcode":0}'
+
+        check_url = ''
+        check_data = ''
+
+        class TPLinkEXClientTest(TPLinkEXClient):
+            def _request(self, url, method='POST', data_str=None, encrypt=False):
+                nonlocal check_url, check_data
+                check_url = url
+                check_data = data_str
+                return 200, response
+
+        client = TPLinkEXClientTest('', '')
+        client.set_vpn(VPN.OPEN_VPN, True)
+
+        self.assertIn('http:///cgi_gdpr?9?_=', check_url)
+        self.assertEqual(check_data, '{"data":{"stack":"0,0,0,0,0,0","pstack":"0,0,0,0,0,0",'
+                                     '"enable":"1"},"operation":"so","oid":"DEV2_OPENVPN"}')
+
 
 if __name__ == '__main__':
     main()
diff --git a/tplinkrouterc6u/client/ex.py b/tplinkrouterc6u/client/ex.py
@@ -4,8 +4,15 @@
 from macaddress import EUI48
 from ipaddress import IPv4Address
 from logging import Logger
-from tplinkrouterc6u.common.package_enum import Connection
-from tplinkrouterc6u.common.dataclass import Firmware, Status, Device, IPv4Reservation, IPv4DHCPLease, IPv4Status
+from tplinkrouterc6u.common.package_enum import Connection, VPN
+from tplinkrouterc6u.common.dataclass import (
+    Firmware,
+    Status,
+    Device,
+    IPv4Reservation,
+    IPv4DHCPLease,
+    IPv4Status,
+    VPNStatus)
 from tplinkrouterc6u.common.exception import ClientException, ClientError
 from tplinkrouterc6u.client.mr import TPLinkMRClientBase
 
@@ -293,3 +300,33 @@ def _req_login(self) -> None:
             if self._logger:
                 self._logger.debug(error)
             raise ClientException(error)
+
+    def get_vpn_status(self) -> VPNStatus:
+        status = VPNStatus()
+        acts = [
+            self.ActItem(self.ActItem.GET, 'DEV2_OPENVPN', attrs=['enable']),
+            self.ActItem(self.ActItem.GET, 'DEV2_PPTPVPN', attrs=['enable']),
+            self.ActItem(self.ActItem.GL, 'DEV2_OVPN_CLIENT', attrs=['connAct']),
+            self.ActItem(self.ActItem.GL, 'DEV2_PVPN_CLIENT', attrs=['connAct']),
+        ]
+        _, values = self.req_act(acts)
+
+        status.openvpn_enable = values[0]['enable'] == '1'
+        status.pptpvpn_enable = values[1]['enable'] == '1'
+
+        for item in values[2]:
+            if item['connAct'] == '1':
+                status.openvpn_clients_total += 1
+
+        for item in values[3]:
+            if item['connAct'] == '1':
+                status.pptpvpn_clients_total += 1
+
+        return status
+
+    def set_vpn(self, vpn: VPN, enable: bool) -> None:
+        acts = [
+            self.ActItem(self.ActItem.SET, "DEV2_" + vpn.value, attrs=[f'"enable":"{int(enable)}"'])
+        ]
+
+        self.req_act(acts)