support for multiple tokens (#309)

* baseapi: support for multiple tokens

To exceed limits of API for single token add support for multiple
tokens used in round-robin style.

* pass multiple tokens to new instances

* nitpicking

* pass internally token list instead of single token

Co-authored-by: Tomáš Čech <sleep_walker@gnu.org>

Author: tc-ubnt

Dockerfile

@@ -1,4 +1,4 @@
-# Dokcerfile made to test this
+# Dockerfile made to test this
 FROM python:alpine
 
 MAINTAINER Lorenzo Setale <lorenzo@setale.me>

digitalocean/Action.py

@@ -7,7 +7,6 @@
 class Action(BaseAPI):
     def __init__(self, *args, **kwargs):
         self.id = None
-        self.token = None
         self.status = None
         self.type = None
         self.started_at = None

digitalocean/Droplet.py

@@ -226,7 +226,7 @@ def _perform_action(self, params, return_dict=True):
             return action
         else:
             action = action[u'action']
-            return_action = Action(token=self.token)
+            return_action = Action(token=self.tokens)
             # Loading attributes
             for attr in action.keys():
                 setattr(return_action, attr, action[attr])
@@ -590,7 +590,7 @@ def get_actions(self):
         actions = []
         for action_dict in answer['actions']:
             action = Action(**action_dict)
-            action.token = self.token
+            action.token = self.tokens
             action.droplet_id = self.id
             action.load()
             actions.append(action)
@@ -603,7 +603,7 @@ def get_action(self, action_id):
             action_id (int): id of action
         """
         return Action.get_object(
-            api_token=self.token,
+            api_token=self.tokens,
             action_id=action_id
         )
 
@@ -616,7 +616,7 @@ def get_snapshots(self):
         for id in self.snapshot_ids:
             snapshot = Image()
             snapshot.id = id
-            snapshot.token = self.token
+            snapshot.token = self.tokens
             snapshots.append(snapshot)
         return snapshots
 
@@ -630,7 +630,7 @@ def get_kernel_available(self):
         while True:
             for jsond in data[u'kernels']:
                 kernel = Kernel(**jsond)
-                kernel.token = self.token
+                kernel.token = self.tokens
                 kernels.append(kernel)
             try:
                 url = data[u'links'][u'pages'].get(u'next')
@@ -653,7 +653,7 @@ def update_volumes_data(self):
         self.volumes = list()
 
         for volume_id in self.volume_ids:
-            volume = Volume().get_object(self.token, volume_id)
+            volume = Volume().get_object(self.tokens, volume_id)
             self.volumes.append(volume)
 
 

digitalocean/Manager.py

@@ -34,7 +34,7 @@ def get_account(self):
         """
             Returns an Account object.
         """
-        return Account.get_object(api_token=self.token)
+        return Account.get_object(api_token=self.tokens)
 
     def get_balance(self):
         """
@@ -50,7 +50,7 @@ def get_all_regions(self):
         regions = list()
         for jsoned in data['regions']:
             region = Region(**jsoned)
-            region.token = self.token
+            region.token = self.tokens
             regions.append(region)
         return regions
 
@@ -69,7 +69,7 @@ def get_all_droplets(self, params=None, tag_name=None):
         droplets = list()
         for jsoned in data['droplets']:
             droplet = Droplet(**jsoned)
-            droplet.token = self.token
+            droplet.token = self.tokens
 
             for net in droplet.networks['v4']:
                 if net['type'] == 'private':
@@ -100,7 +100,7 @@ def get_droplet(self, droplet_id):
         """
             Return a Droplet by its ID.
         """
-        return Droplet.get_object(api_token=self.token, droplet_id=droplet_id)
+        return Droplet.get_object(api_token=self.tokens, droplet_id=droplet_id)
 
     def get_all_sizes(self):
         """
@@ -110,7 +110,7 @@ def get_all_sizes(self):
         sizes = list()
         for jsoned in data['sizes']:
             size = Size(**jsoned)
-            size.token = self.token
+            size.token = self.tokens
             sizes.append(size)
         return sizes
 
@@ -127,7 +127,7 @@ def get_images(self, private=False, type=None):
         images = list()
         for jsoned in data['images']:
             image = Image(**jsoned)
-            image.token = self.token
+            image.token = self.tokens
             images.append(image)
         return images
 
@@ -144,7 +144,7 @@ def get_image(self, image_id_or_slug):
             Return a Image by its ID/Slug.
         """
         return Image.get_object(
-            api_token=self.token,
+            api_token=self.tokens,
             image_id_or_slug=image_id_or_slug,
         )
 
@@ -166,7 +166,7 @@ def get_global_images(self):
         images = list()
         for i in data:
             if i.public:
-                i.token = self.token
+                i.token = self.tokens
                 images.append(i)
         return images
 
@@ -194,15 +194,15 @@ def get_all_domains(self):
         domains = list()
         for jsoned in data['domains']:
             domain = Domain(**jsoned)
-            domain.token = self.token
+            domain.token = self.tokens
             domains.append(domain)
         return domains
 
     def get_domain(self, domain_name):
         """
             Return a Domain by its domain_name
         """
-        return Domain.get_object(api_token=self.token, domain_name=domain_name)
+        return Domain.get_object(api_token=self.tokens, domain_name=domain_name)
 
     def get_all_sshkeys(self):
         """
@@ -212,15 +212,15 @@ def get_all_sshkeys(self):
         ssh_keys = list()
         for jsoned in data['ssh_keys']:
             ssh_key = SSHKey(**jsoned)
-            ssh_key.token = self.token
+            ssh_key.token = self.tokens
             ssh_keys.append(ssh_key)
         return ssh_keys
 
     def get_ssh_key(self, ssh_key_id):
         """
             Return a SSHKey object by its ID.
         """
-        return SSHKey.get_object(api_token=self.token, ssh_key_id=ssh_key_id)
+        return SSHKey.get_object(api_token=self.tokens, ssh_key_id=ssh_key_id)
 
     def get_all_tags(self):
         """
@@ -235,7 +235,7 @@ def get_action(self, action_id):
         """
             Return an Action object by a specific ID.
         """
-        return Action.get_object(api_token=self.token, action_id=action_id)
+        return Action.get_object(api_token=self.tokens, action_id=action_id)
 
     def get_all_floating_ips(self):
         """
@@ -245,15 +245,15 @@ def get_all_floating_ips(self):
         floating_ips = list()
         for jsoned in data['floating_ips']:
             floating_ip = FloatingIP(**jsoned)
-            floating_ip.token = self.token
+            floating_ip.token = self.tokens
             floating_ips.append(floating_ip)
         return floating_ips
 
     def get_floating_ip(self, ip):
         """
             Returns a of FloatingIP object by its IP address.
         """
-        return FloatingIP.get_object(api_token=self.token, ip=ip)
+        return FloatingIP.get_object(api_token=self.tokens, ip=ip)
 
     def get_all_load_balancers(self):
         """
@@ -264,7 +264,7 @@ def get_all_load_balancers(self):
         load_balancers = list()
         for jsoned in data['load_balancers']:
             load_balancer = LoadBalancer(**jsoned)
-            load_balancer.token = self.token
+            load_balancer.token = self.tokens
             load_balancer.health_check = HealthCheck(**jsoned['health_check'])
             load_balancer.sticky_sessions = StickySessions(**jsoned['sticky_sessions'])
             forwarding_rules = list()
@@ -281,7 +281,7 @@ def get_load_balancer(self, id):
             Args:
                 id (str): Load Balancer ID
         """
-        return LoadBalancer.get_object(api_token=self.token, id=id)
+        return LoadBalancer.get_object(api_token=self.tokens, id=id)
 
     def get_certificate(self, id):
         """
@@ -290,7 +290,7 @@ def get_certificate(self, id):
             Args:
                 id (str): Certificate ID
         """
-        return Certificate.get_object(api_token=self.token, cert_id=id)
+        return Certificate.get_object(api_token=self.tokens, cert_id=id)
 
     def get_all_certificates(self):
         """
@@ -300,7 +300,7 @@ def get_all_certificates(self):
         certificates = list()
         for jsoned in data['certificates']:
             cert = Certificate(**jsoned)
-            cert.token = self.token
+            cert.token = self.tokens
             certificates.append(cert)
 
         return certificates
@@ -310,7 +310,7 @@ def get_snapshot(self, snapshot_id):
             Return a Snapshot by its ID.
         """
         return Snapshot.get_object(
-            api_token=self.token, snapshot_id=snapshot_id
+            api_token=self.tokens, snapshot_id=snapshot_id
         )
 
     def get_all_snapshots(self):
@@ -319,7 +319,7 @@ def get_all_snapshots(self):
         """
         data = self.get_data("snapshots/")
         return [
-            Snapshot(token=self.token, **snapshot)
+            Snapshot(token=self.tokens, **snapshot)
             for snapshot in data['snapshots']
         ]
 
@@ -329,7 +329,7 @@ def get_droplet_snapshots(self):
         """
         data = self.get_data("snapshots?resource_type=droplet")
         return [
-            Snapshot(token=self.token, **snapshot)
+            Snapshot(token=self.tokens, **snapshot)
             for snapshot in data['snapshots']
         ]
 
@@ -339,7 +339,7 @@ def get_volume_snapshots(self):
         """
         data = self.get_data("snapshots?resource_type=volume")
         return [
-            Snapshot(token=self.token, **snapshot)
+            Snapshot(token=self.tokens, **snapshot)
             for snapshot in data['snapshots']
         ]
 
@@ -355,15 +355,15 @@ def get_all_volumes(self, region=None):
         volumes = list()
         for jsoned in data['volumes']:
             volume = Volume(**jsoned)
-            volume.token = self.token
+            volume.token = self.tokens
             volumes.append(volume)
         return volumes
 
     def get_volume(self, volume_id):
         """
             Returns a Volume object by its ID.
         """
-        return Volume.get_object(api_token=self.token, volume_id=volume_id)
+        return Volume.get_object(api_token=self.tokens, volume_id=volume_id)
 
     def get_all_projects(self):
         """
@@ -403,7 +403,7 @@ def get_all_firewalls(self):
         firewalls = list()
         for jsoned in data['firewalls']:
             firewall = Firewall(**jsoned)
-            firewall.token = self.token
+            firewall.token = self.tokens
             in_rules = list()
             for rule in jsoned['inbound_rules']:
                 in_rules.append(InboundRule(**rule))
@@ -420,7 +420,7 @@ def get_firewall(self, firewall_id):
             Return a Firewall by its ID.
         """
         return Firewall.get_object(
-            api_token=self.token,
+            api_token=self.tokens,
             firewall_id=firewall_id,
         )
 

digitalocean/Volume.py

@@ -189,7 +189,7 @@ def get_snapshots(self):
         snapshots = list()
         for jsond in data[u'snapshots']:
             snapshot = Snapshot(**jsond)
-            snapshot.token = self.token
+            snapshot.token = self.tokens
             snapshots.append(snapshot)
 
         return snapshots

digitalocean/baseapi.py

@@ -47,7 +47,8 @@ class BaseAPI(object):
     """
         Basic api class for
     """
-    token = ""
+    tokens = []
+    _last_used = 0
     end_point = "https://api.digitalocean.com/v2/"
 
     def __init__(self, *args, **kwargs):
@@ -59,7 +60,7 @@ def __init__(self, *args, **kwargs):
 
         for attr in kwargs.keys():
             setattr(self, attr, kwargs[attr])
-        
+
         parsed_url = urlparse.urlparse(self.end_point)
         if not parsed_url.scheme or not parsed_url.netloc:
             raise EndPointError("Provided end point is not a valid URL. Please use a valid URL")
@@ -123,7 +124,9 @@ def __perform_request(self, url, type=GET, params=None):
             kwargs['timeout'] = timeout
 
         # remove token from log
-        headers_str = str(headers).replace(self.token.strip(), 'TOKEN')
+        headers_str = str(headers)
+        for i, token in enumerate(self.tokens):
+            headers_str = headers_str.replace(token.strip(), 'TOKEN%s' % i)
         self._log.debug('%s %s %s:%s %s %s' %
                         (type, url, payload, params, headers_str, timeout))
 
@@ -162,6 +165,23 @@ def __init_ratelimit(self, headers):
         # Add the account requests limit reset time
         self.ratelimit_reset = headers.get('Ratelimit-Reset', None)
 
+    @property
+    def token(self):
+        # use all the tokens round-robin style
+        if self.tokens:
+            self._last_used = (self._last_used + 1) % len(self.tokens)
+            return self.tokens[self._last_used]
+        return ""
+
+    @token.setter
+    def token(self, token):
+        self._last_used = 0
+        if isinstance(token, list):
+            self.tokens = token
+        else:
+            # for backward compatibility
+            self.tokens = [token]
+
     def get_timeout(self):
         """
             Checks if any timeout for the requests to DigitalOcean is required.