meaksh · September 4, 2019 14:18 · Sep 4, 2019 · Sep 4, 2019
diff --git a/public_clouds.py → public_clouds_grain.py b/public_clouds.py → public_clouds_grain.py
diff --git a/public_clouds.py b/public_clouds.py
@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-                                                                                                                                                      
+'''
+    This grain module is only loaded in case of a public cloud instance.
+
+    Supported Instances: AWS EC2, Azure and Google Compute Engine instances
+
+    Returns a grain called "instance_id" containing the virtual instance ID
+    according to the Public Cloud provider. The data is gathered using the
+    internal API available from within the instance.
+
+    Author: Pablo Suárez Hernández <[email protected]>
+    Based on: https://docs.saltstack.com/en/latest/ref/grains/all/salt.grains.metadata.html
+
+'''
+from __future__ import absolute_import, print_function, unicode_literals
+
+# Import python libs
+import os
+import socket
+
+# Import salt libs
+import salt.utils.http as http
+
+# Internal instance information
+INTERNAL_API_IP = '169.254.169.254'
+HOST = 'http://{0}/'.format(INTERNAL_API_IP)
+
+IS_AMAZON = False
+IS_AZURE = False
+IS_GOOGLE = False
+
+AMAZON_URL_PATH = 'latest/meta-data/instance-id'
+AZURE_URL_PATH = 'metadata/instance/compute/vmId?api-version=2017-08-01&format=text'
+GOOGLE_URL_PATH = 'computeMetadata/v1/instance/id'
+
+
+def __virtual__():
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.settimeout(0.1)
+    result = sock.connect_ex((INTERNAL_API_IP, 80))
+    if result != 0:
+        return False
+    if http.query(os.path.join(HOST, AMAZON_URL_PATH), status=True).get('status') == 200:
+        global IS_AMAZON
+        IS_AMAZON = True
+        return True
+    elif http.query(os.path.join(HOST, AZURE_URL_PATH), status=True, header_dict={"Metadata":"true"}).get('status') == 200:
+        global IS_AZURE
+        IS_AZURE = True
+        return True
+    elif http.query(os.path.join(HOST, GOOGLE_URL_PATH), status=True, header_dict={"Metadata-Flavor": "Google"}).get('status') == 200:
+        global IS_GOOGLE
+        IS_GOOGLE = True
+        return True
+    return False
+
+
+def instance_id():
+    ret = {}
+    if IS_AMAZON:
+        ret['instance_id'] = http.query(HOST + AMAZON_URL_PATH)['body']
+    elif IS_AZURE:
+        ret['instance_id'] = http.query(HOST + AZURE_URL_PATH,
+                                        header_dict={"Metadata":"true"})['body']
+    elif IS_GOOGLE:
+        ret['instance_id'] = http.query(HOST + GOOGLE_URL_PATH,
+                                        header_dict={"Metadata-Flavor": "Google"})['body']
+    return ret