Browse code

Merge pull request #1118 from Patouche/issues/1117

[#1117] Improve speed when running on EC2 instance

Florent Viard authored on 2020/06/27 21:06:44
Showing 1 changed files
... ...
@@ -17,6 +17,26 @@ import sys
17 17
 import json
18 18
 from . import Progress
19 19
 from .SortedDict import SortedDict
20
+import datetime
21
+from .ExitCodes import EX_OSFILE
22
+
23
+try:
24
+    import dateutil.parser
25
+except ImportError:
26
+    sys.stderr.write(u"""
27
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
28
+ImportError trying to import dateutil.parser.
29
+Please install the python dateutil module:
30
+$ sudo apt-get install python-dateutil
31
+  or
32
+$ sudo yum install python-dateutil
33
+  or
34
+$ pip install python-dateutil
35
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
36
+""")
37
+    sys.stderr.flush()
38
+    sys.exit(EX_OSFILE)
39
+
20 40
 try:
21 41
     # python 3 support
22 42
     import httplib
... ...
@@ -54,6 +74,12 @@ def config_unicodise(string, encoding = "utf-8", errors = "replace"):
54 54
     except UnicodeDecodeError:
55 55
         raise UnicodeDecodeError("Conversion to unicode failed: %r" % string)
56 56
 
57
+def config_date_to_python(date):
58
+    """
59
+    Convert a string formated like '2020-06-27T15:56:34Z' into a python datetime 
60
+    """
61
+    return dateutil.parser.parse(date, fuzzy=True)
62
+
57 63
 def is_bool_true(value):
58 64
     """Check to see if a string is true, yes, on, or 1
59 65
 
... ...
@@ -94,6 +120,8 @@ class Config(object):
94 94
     secret_key = u""
95 95
     access_token = u""
96 96
     _access_token_refresh = True
97
+    _access_token_expiration = None
98
+    _access_token_last_update = None
97 99
     host_base = u"s3.amazonaws.com"
98 100
     host_bucket = u"%(bucket)s.s3.amazonaws.com"
99 101
     kms_key = u""    #can't set this and Server Side Encryption at the same time
... ...
@@ -273,7 +301,7 @@ class Config(object):
273 273
         Get credentials from IAM authentication
274 274
         """
275 275
         try:
276
-            conn = httplib.HTTPConnection(host='169.254.169.254', timeout = 2)
276
+            conn = httplib.HTTPConnection(host='169.254.169.254', timeout=2)
277 277
             conn.request('GET', "/latest/meta-data/iam/security-credentials/")
278 278
             resp = conn.getresponse()
279 279
             files = resp.read()
... ...
@@ -286,6 +314,11 @@ class Config(object):
286 286
                     Config().update_option('access_key', config_unicodise(creds['AccessKeyId']))
287 287
                     Config().update_option('secret_key', config_unicodise(creds['SecretAccessKey']))
288 288
                     Config().update_option('access_token', config_unicodise(creds['Token']))
289
+                    expiration = config_date_to_python(config_unicodise(creds['Expiration']))
290
+                    # Add a timedelta to prevent any expiration if the EC2 machine is not at the right date
291
+                    self._access_token_expiration = expiration - datetime.timedelta(minutes=15)
292
+                    self._access_token_last_update = config_date_to_python(config_unicodise(creds['LastUpdated']))
293
+                    # Others variables : Code / Type
289 294
                 else:
290 295
                     raise IOError
291 296
             else:
... ...
@@ -295,6 +328,13 @@ class Config(object):
295 295
 
296 296
     def role_refresh(self):
297 297
         if self._access_token_refresh:
298
+            now = datetime.datetime.now(datetime.timezone.utc)
299
+            if self._access_token_expiration \
300
+               and now < self._access_token_expiration \
301
+               and self._access_token_last_update \
302
+               and self._access_token_last_update <= now:
303
+                # current token is still valid. No need to refresh it
304
+                return
298 305
             try:
299 306
                 self.role_config()
300 307
             except Exception: