1 files changed, 63 insertions, 46 deletions

diff --git a/wrappers/python/src/server.py b/wrappers/python/src/server.py
index b765ede..dce5d51 100644
--- a/wrappers/python/src/server.py
+++ b/wrappers/python/src/server.py
@@ -1,5 +1,6 @@
-from . import lib, cServers, cServerLocations
-from ctypes import cast, POINTER
+from . import lib, cServer, cServers, cServerLocations, cServerProfiles
+from ctypes import cast, POINTER, c_char_p
+from datetime import datetime
 
 
 class Profile:
@@ -9,61 +10,85 @@ class Profile:
         self.default_gateway = default_gateway
 
     def __str__(self):
-        return f"Profile: {self.display_name}"
+        return self.display_name
+
+
+class Profiles:
+    def __init__(self, profiles, current):
+        self.profiles = profiles
+        self.current_index = current
+
+    @property
+    def current(self):
+        if self.current_index < len(self.profiles):
+            return self.profiles[self.current_index]
+        return None
 
 
 class Server:
-    def __init__(self, url, display_name, profiles, current_profile, expire_time):
+    def __init__(self, url, display_name, profiles=None, expire_time=0):
         self.url = url
         self.display_name = display_name
         self.profiles = profiles
         self.current_profile = None
-        if current_profile < len(profiles):
-            self.current_profile = profiles[current_profile]
-        self.expire_time = expire_time
+        self.expire_time = datetime.fromtimestamp(expire_time)
 
     def __str__(self):
-        return f"Server: {self.url}, with current profile: {self.current_profile}"
+        return self.display_name
 
 
 class InstituteServer(Server):
-    def __init__(
-        self, url, display_name, support_contact, profiles, current_profile, expire_time
-    ):
-        super().__init__(url, display_name, profiles, current_profile, expire_time)
+    def __init__(self, url, display_name, support_contact, profiles, expire_time):
+        super().__init__(url, display_name, profiles, expire_time)
         self.support_contact = support_contact
 
-    def __str__(self):
-        return f"Institute Server: {self.display_name}"
-
 
 class SecureInternetServer(Server):
     def __init__(
         self,
-        url,
+        org_id,
         display_name,
         support_contact,
         profiles,
-        current_profile,
         expire_time,
         country_code,
     ):
-        super().__init__(url, display_name, profiles, current_profile, expire_time)
+        super().__init__(org_id, display_name, profiles, expire_time)
+        self.org_id = org_id
         self.support_contact = support_contact
         self.country_code = country_code
 
-    def __str__(self):
-        return f"Secure Internet Server: {self.display_name} with country {self.country_code}"
-
 
 def get_type_for_str(type_str: str):
-    if type_str is "secure_internet":
+    if type_str == "secure_internet":
         return SecureInternetServer
-    if type_str is "custom_server":
+    if type_str == "custom_server":
         return Server
     return InstituteServer
 
 
+def get_profiles(ptr):
+    if not ptr:
+        return []
+    profiles = []
+    _profiles = ptr.contents
+    current_profile = _profiles.current
+    if not _profiles.profiles:
+        return []
+    for i in range(_profiles.total_profiles):
+        if not _profiles.profiles[i]:
+            continue
+        profile = _profiles.profiles[i].contents
+        profiles.append(
+            Profile(
+                profile.identifier.decode("utf-8"),
+                profile.display_name.decode("utf-8"),
+                profile.default_gateway == 1,
+            )
+        )
+    return Profiles(profiles, current_profile)
+
+
 def get_server(ptr, _type=None):
     if not ptr:
         return None
@@ -79,31 +104,13 @@ def get_server(ptr, _type=None):
         support_contact = []
         for i in range(current_server.total_support_contact):
             support_contact.append(current_server.support_contact[i].decode("utf-8"))
-    profiles = []
-    if not current_server.profiles:
-        return None
-
-    _profiles = current_server.profiles.contents
-    current_profile = _profiles.current
-    for i in range(_profiles.total_profiles):
-        if not _profiles.profiles or not _profiles.profiles[i]:
-            return None
-        profile = _profiles.profiles[i].contents
-        profiles.append(
-            Profile(
-                profile.identifier.decode("utf-8"),
-                profile.display_name.decode("utf-8"),
-                profile.default_gateway == 1,
-            )
-        )
-
+    profiles = get_profiles(current_server.profiles)
     if _type is SecureInternetServer:
         return SecureInternetServer(
             identifier,
             display_name,
             support_contact,
             profiles,
-            current_profile,
             current_server.expire_time,
             current_server.country_code.decode("utf-8"),
         )
@@ -113,12 +120,21 @@ def get_server(ptr, _type=None):
             display_name,
             support_contact,
             profiles,
-            current_profile,
             current_server.expire_time,
         )
-    return Server(
-        identifier, display_name, profiles, current_profile, current_server.expire_time
-    )
+    return Server(identifier, display_name, profiles, current_server.expire_time)
+
+
+def get_transition_server(ptr):
+    server = get_server(cast(ptr, POINTER(cServer)))
+    lib.FreeServer(ptr)
+    return server
+
+
+def get_transition_profiles(ptr):
+    profiles = get_profiles(cast(ptr, POINTER(cServerProfiles)))
+    lib.FreeProfiles(ptr)
+    return profiles
 
 
 def get_servers(ptr):
@@ -147,6 +163,7 @@ def get_servers(ptr):
         return returned
     return None
 
+
 def get_locations(ptr):
     if ptr:
         locations = cast(ptr, POINTER(cServerLocations)).contents