From 80603d13d303713f0dae19f79ffe2943c52c09a7 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Fri, 8 Nov 2019 10:25:06 +0200
Subject: [PATCH] cloudinfo: Add Azure metadata support

---
 pupy/packages/all/cloudinfo.py | 148 +++++++++++++++++++++++----------
 1 file changed, 106 insertions(+), 42 deletions(-)

diff --git a/pupy/packages/all/cloudinfo.py b/pupy/packages/all/cloudinfo.py
index 197815fb..9dc23b28 100644
--- a/pupy/packages/all/cloudinfo.py
+++ b/pupy/packages/all/cloudinfo.py
@@ -8,32 +8,87 @@ import string
 
 is_dict = re.compile(r'^(\d+)=([^/]+)$')
 
-def get(path, version='latest', section='meta-data'):
-    if path == '/':
-        path = ''
+
+METADATA_SERVER = 'http://169.254.169.254'
+METADATA_PROVIDER = None
+METADATA_ROOT = None
+METADATA_HEADERS = []
+METADATA_PARAMS = ''
+
+
+def _probe():
+    global METADATA_ROOT
+    global METADATA_PROVIDER
+    global METADATA_HEADERS
+    global METADATA_PARAMS
+
+    if get('', 'latest/meta-data'):
+        METADATA_PROVIDER = 'EC2'
+        METADATA_ROOT = 'latest'
+
+    elif get('', 'metadata/v1'):
+        METADATA_PROVIDER = 'DO'
+        METADATA_ROOT = 'metadata/v1'
+
+    else:
+        result = get('', 'metadata', [('Metadata', 'true')], 400)
+        if not result:
+            raise ValueError('Unknown metadata implementation')
+
+        error = json.loads(result)
+        version = error['newest-versions'][0]
+
+        METADATA_PROVIDER = 'MS'
+        METADATA_HEADERS = [('Metadata', 'true')]
+        METADATA_ROOT = 'metadata/instance'
+        METADATA_PARAMS = '?api-version={}&format=json'.format(version)
+
+    return METADATA_PROVIDER, METADATA_ROOT, METADATA_HEADERS
+
+
+def get(path, root=None, headers=None, code=200):
+    opener = urllib2.build_opener()
+    opener.addheaders = \
+        METADATA_HEADERS if headers is None else headers
+
+    path = '/'.join([x for x in path.split('/') if x])
+
+    uri = '/'.join([
+        METADATA_SERVER,
+        METADATA_ROOT if root is None else root,
+        path or ''
+    ]) + METADATA_PARAMS
 
     try:
-        return urllib2.urlopen('http://169.254.169.254/{}/{}/{}'.format(
-            version, section, path), timeout=2
-        ).read()
+        response = opener.open(uri, timeout=2)
+        if response.code != code:
+            return None
 
-    except:
-        return None
+    except urllib2.HTTPError as e:
+        if e.code == code:
+            return e.fp.read()
+        else:
+            return None
 
-def dir(path, version='latest', section='meta-data'):
-    path = '/'.join([x for x in path.split('/') if x]) + '/'
-    data = get(path, version, section)
+    return response.read()
+
+
+def metadir(path):
+    data = get(path)
 
     if data:
         return [x for x in data.split('\n') if x]
     else:
         return data
 
-def list(path, version='latest', section='meta-data'):
-    result = {}
-    for item in dir(path, version, section):
-        result[item] = get('/'.join(path, item))
-    return result
+
+def isint(x):
+    try:
+        int(x)
+        return True
+    except:
+        return False
+
 
 def valueconv(x):
     try:
@@ -59,7 +114,7 @@ def valueconv(x):
     except:
         pass
 
-    if x == 'none':
+    if x is None or x == 'none':
         return None
     elif x.endswith('\n'):
         return x
@@ -68,54 +123,63 @@ def valueconv(x):
 
     return x
 
-def as_dict(path='', version='latest', section='meta-data'):
+
+def as_dict(path=''):
     result = {}
-    dirs = dir(path, version, section)
+    dirs = metadir(path)
     if not dirs:
         return None
 
     for item in dirs:
         if item.endswith('/'):
-            records = as_dict(path+item, version, section)
+            records = as_dict('/'.join([path, item]))
             if records:
                 result[item[:-1]] = records
 
         elif is_dict.match(item):
             idx, name = is_dict.match(item).groups()
-            records = as_dict(path+idx+'/', version, section)
+            records = as_dict('/'.join([path, idx + '/']))
             if records:
                 result[name] = records
         else:
-            result[item] = valueconv(get(path+item, version, section))
+            result[item] = valueconv(get('/'.join([path, item])))
+
+    if isinstance(result, dict) and all(
+            (isint(key) and int(key) < len(result)) for key in result):
+        as_list = [None] * len(result)
+        for key, value in result.iteritems():
+            as_list[int(key)] = value
+
+        result = as_list
 
     return result
 
+
 def metadata():
-    result = as_dict()
-    if result:
+    if METADATA_PROVIDER is None:
+        try:
+            _probe()
+        except ValueError:
+            return None
+
+    if METADATA_PROVIDER == 'EC2':
         result = {
-            'meta-data': result
+            'meta-data': as_dict('meta-data'),
+            'dynamic': as_dict('dynamic'),
         }
 
-        user_data = get('', section='user-data')
+        user_data = as_dict('user-data')
         if user_data:
             try:
-                result.update({
-                    'user-data': dict(x.split('=', 1) for x in user_data.split(';'))
-                })
+                result['user-data'] = dict(
+                    x.split('=', 1) for x in user_data.split(';')
+                )
             except:
-                result.update({
-                    'user-data': user_data
-                })
+                result['user-data'] = user_data
 
-        result.update({
-            'dynamic': as_dict('/', section='dynamic')
-        })
-        return 'EC2', result
+    elif METADATA_PROVIDER == 'MS':
+        result = json.loads(get(''))
+    else:
+        result = as_dict()
 
-    result = as_dict('', 'metadata', 'v1')
-
-    if result:
-        return 'DO', result
-
-    return None, None
+    return METADATA_PROVIDER, result