get_ip, get_public_ip and get_private_ip

Author: elnygren

upcloud_api/server.py

@@ -1,10 +1,9 @@
 from __future__ import unicode_literals
 from __future__ import absolute_import
 
-import itertools
 from time import sleep
 
-from upcloud_api import Storage, IPAddress, UpCloudAPIError, OperatingSystems
+from upcloud_api import Storage, IPAddress, OperatingSystems
 from upcloud_api.utils import try_it_n_times
 
 
@@ -392,48 +391,46 @@ def to_dict(self):
         del fields['cloud_manager']
         return fields
 
-    def get_public_ip(self, addr_family='IPv4', strict=False):
+    def get_ip(self, access='public', addr_family=None, strict=None):
         """
-        Return a server's public IP.
+        Return the server's IP address.
 
         Params:
-        - addr_family: prefer IPv4 (default) or IPv6.
-        - strict mode (false/off by default): only return IP if it belongs to
-                                              addr_family (IPv4 or IPv6).
-
-        Tries to fetch Server data from API if ip_addresses not set.
-
-        New in 3.4:
-        - possibility to specify which protocol is preferred via addr_family,
-          instead of always preferring IPv4.
-        - strict mode
+        - addr_family: IPv4, IPv6 or None. None prefers IPv4 but will
+                       return IPv6 if IPv4 addr was not available.
+        - access: 'public' or 'private'
         """
-        if addr_family not in ['IPv4', 'IPv6']:
-            raise Exception("`addr_family` must be 'IPv4' or 'IPv6'")
+        if addr_family not in ['IPv4', 'IPv6', None]:
+            raise Exception("`addr_family` must be 'IPv4', 'IPv6' or None")
+
+        if access not in ['private', 'public']:
+            raise Exception("`access` must be 'public' or 'private'")
 
         if not hasattr(self, 'ip_addresses'):
             self.populate()
 
-        # server can have several public IPs
-        public_ip_addrs = []
-        for ip_addr in self.ip_addresses:
-            if ip_addr.access == 'public':
-                public_ip_addrs.append(ip_addr)
-
-        if not public_ip_addrs:
-            return None
+        # server can have several public or private IPs
+        ip_addrs = [
+            ip_addr for ip_addr in self.ip_addresses
+            if ip_addr.access == access
+        ]
 
-        # prefer addr_family
-        for ip_addr in public_ip_addrs:
-            if ip_addr.family == addr_family:
+        # prefer addr_family (or IPv4 if none given)
+        preferred_family = addr_family if addr_family else 'IPv4'
+        for ip_addr in ip_addrs:
+            if ip_addr.family == preferred_family:
                 return ip_addr.address
 
-        # strict mode: either find addr_family or don't
-        if strict:
-            return None
+        # any IP (of the right access) will do if available and addr_family is None
+        return ip_addrs[0].address if ip_addrs and not addr_family else None
+
+    def get_public_ip(self, addr_family=None, *args, **kwargs):
+        """Alias for get_ip('public')"""
+        return self.get_ip('public', addr_family, *args, **kwargs)
 
-        # not stict mode: any public IP will do if addr_family didn't match
-        return public_ip_addrs[0].address
+    def get_private_ip(self, addr_family=None, *args, **kwargs):
+        """Alias for get_ip('private')"""
+        return self.get_ip('private', addr_family, *args, **kwargs)
 
     def _wait_for_state_change(self, target_states, update_interval=10):
         """