Exports + Python: Use an enum for server type

2 files changed, 19 insertions, 11 deletions

diff --git a/wrappers/python/eduvpn_common/loader.py b/wrappers/python/eduvpn_common/loader.py
index c5709fb..e57e81d 100644
--- a/wrappers/python/eduvpn_common/loader.py
+++ b/wrappers/python/eduvpn_common/loader.py
@@ -64,18 +64,18 @@ def initialize_functions(lib: CDLL) -> None:
     lib.DiscoOrganizations.argtypes, lib.DiscoOrganizations.restype = [], DataError
     lib.DiscoServers.argtypes, lib.DiscoServers.restype = [], DataError
     lib.GetConfig.argtypes, lib.GetConfig.restype = [
-        c_char_p,
+        c_int,
         c_char_p,
         c_int,
         c_char_p,
     ], DataError
     lib.AddServer.argtypes, lib.AddServer.restype = [
-        c_char_p,
+        c_int,
         c_char_p,
     ], c_char_p
     lib.CurrentServer.argtypes, lib.CurrentServer.restype = [], DataError
     lib.RemoveServer.argtypes, lib.RemoveServer.restype = [
-        c_char_p,
+        c_int,
         c_char_p,
     ], c_char_p
     lib.ServerList.argtypes, lib.ServerList.restype = [], DataError
diff --git a/wrappers/python/eduvpn_common/main.py b/wrappers/python/eduvpn_common/main.py
index f22e4d3..87de430 100644
--- a/wrappers/python/eduvpn_common/main.py
+++ b/wrappers/python/eduvpn_common/main.py
@@ -1,3 +1,4 @@
+from enum import IntEnum
 from typing import Any, Callable, Iterator, Optional
 
 from eduvpn_common.loader import initialize_functions, load_lib
@@ -16,6 +17,13 @@ def forwardError(error: bytes | str):
     raise WrappedError(error.decode("utf-8"))
 
 
+class ServerType(IntEnum):
+    UNKNOWN = 0
+    INSTITUTE_ACCESS = 1
+    SECURE_INTERNET = 2
+    CUSTOM = 3
+
+
 class EduVPN(object):
     """The main class used to communicate with the Go library.
     It registers the client with the library and then calls the needed appropriate functions
@@ -93,15 +101,15 @@ class EduVPN(object):
         if register_err:
             forwardError(register_err)
 
-    def add_server(self, _type: str, _id: str) -> None:
+    def add_server(self, _type: ServerType, _id: str) -> None:
         """Add a server
 
-        :param _type: str: The type of server: "institute_access", "secure_internet" or "custom_server"
+        :param _type: ServerType: The type of server e.g. SERVER.INSTITUTE_ACCESS
         :param _id: str: The identifier of the server, e.g. "https://vpn.example.com/"
 
         :raises WrappedError: An error by the Go library
         """
-        add_err = self.go_function(self.lib.AddServer, _type, _id)
+        add_err = self.go_function(self.lib.AddServer, int(_type), _id)
 
         if add_err:
             forwardError(add_err)
@@ -140,25 +148,25 @@ class EduVPN(object):
             forwardError(servers_err)
         return servers
 
-    def remove_server(self, _type: str, _id: str) -> None:
+    def remove_server(self, _type: ServerType, _id: str) -> None:
         """Remove a server
 
-        :param _type: str: The type of server: "institute_access", "secure_internet" or "custom_server"
+        :param _type: ServerType: The type of server e.g. SERVER.INSTITUTE_ACCESS
         :param _id: str: The identifier of the server, e.g. "https://vpn.example.com/"
 
         :raises WrappedError: An error by the Go library
         """
-        remove_err = self.go_function(self.lib.RemoveServer, _type, _id)
+        remove_err = self.go_function(self.lib.RemoveServer, int(_type), _id)
 
         if remove_err:
             forwardError(remove_err)
 
     def get_config(
-        self, _type: str, identifier: str, prefer_tcp: bool = False, tokens: str = "{}"
+        self, _type: ServerType, identifier: str, prefer_tcp: bool = False, tokens: str = "{}"
     ) -> Optional[str]:
         """Get an OpenVPN/WireGuard configuration from the server
 
-        :param _type: str: The type of server: "institute_access", "secure_internet" or "custom_server"
+        :param _type: ServerType: The type of server e.g. SERVER.INSTITUTE_ACCESS
         :param identifier: str: The identifier of the server, e.g. URL or ORG ID
         :param prefer_tcp: bool:  (Default value = False): Whether or not to prefer TCP
         :param tokens: str  (Defualt value = ""): The OAuth tokens if available