tools: Add cloud-id command line utility

Add a quick cloud lookup utility in order to more easily determine
the cloud on which an instance is running.

The utility parses standardized attributes from
/run/cloud-init/instance-data.json to print the canonical cloud-id
for the instance. It uses known region maps if necessary to determine
on which specific cloud the instance is running.

Examples:
aws, aws-gov, aws-china, rackspace, azure-china, lxd, openstack, unknown

5 files changed, 319 insertions, 2 deletions

diff --git a/cloudinit/cmd/cloud_id.py b/cloudinit/cmd/cloud_id.py
new file mode 100755
index 00000000..97608921
--- /dev/null
+++ b/cloudinit/cmd/cloud_id.py
@@ -0,0 +1,90 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+"""Commandline utility to list the canonical cloud-id for an instance."""
+
+import argparse
+import json
+import sys
+
+from cloudinit.sources import (
+    INSTANCE_JSON_FILE, METADATA_UNKNOWN, canonical_cloud_id)
+
+DEFAULT_INSTANCE_JSON = '/run/cloud-init/%s' % INSTANCE_JSON_FILE
+
+NAME = 'cloud-id'
+
+
+def get_parser(parser=None):
+    """Build or extend an arg parser for the cloud-id utility.
+
+    @param parser: Optional existing ArgumentParser instance representing the
+        query subcommand which will be extended to support the args of
+        this utility.
+
+    @returns: ArgumentParser with proper argument configuration.
+    """
+    if not parser:
+        parser = argparse.ArgumentParser(
+            prog=NAME,
+            description='Report the canonical cloud-id for this instance')
+    parser.add_argument(
+        '-j', '--json', action='store_true', default=False,
+        help='Report all standardized cloud-id information as json.')
+    parser.add_argument(
+        '-l', '--long', action='store_true', default=False,
+        help='Report extended cloud-id information as tab-delimited string.')
+    parser.add_argument(
+        '-i', '--instance-data', type=str, default=DEFAULT_INSTANCE_JSON,
+        help=('Path to instance-data.json file. Default is %s' %
+              DEFAULT_INSTANCE_JSON))
+    return parser
+
+
+def error(msg):
+    sys.stderr.write('ERROR: %s\n' % msg)
+    return 1
+
+
+def handle_args(name, args):
+    """Handle calls to 'cloud-id' cli.
+
+    Print the canonical cloud-id on which the instance is running.
+
+    @return: 0 on success, 1 otherwise.
+    """
+    try:
+        instance_data = json.load(open(args.instance_data))
+    except IOError:
+        return error(
+            "File not found '%s'. Provide a path to instance data json file"
+            ' using --instance-data' % args.instance_data)
+    except ValueError as e:
+        return error(
+            "File '%s' is not valid json. %s" % (args.instance_data, e))
+    v1 = instance_data.get('v1', {})
+    cloud_id = canonical_cloud_id(
+        v1.get('cloud_name', METADATA_UNKNOWN),
+        v1.get('region', METADATA_UNKNOWN),
+        v1.get('platform', METADATA_UNKNOWN))
+    if args.json:
+        v1['cloud_id'] = cloud_id
+        response = json.dumps(   # Pretty, sorted json
+            v1, indent=1, sort_keys=True, separators=(',', ': '))
+    elif args.long:
+        response = '%s\t%s' % (cloud_id, v1.get('region', METADATA_UNKNOWN))
+    else:
+        response = cloud_id
+    sys.stdout.write('%s\n' % response)
+    return 0
+
+
+def main():
+    """Tool to query specific instance-data values."""
+    parser = get_parser()
+    sys.exit(handle_args(NAME, parser.parse_args()))
+
+
+if __name__ == '__main__':
+    main()
+
+# vi: ts=4 expandtab
diff --git a/cloudinit/cmd/tests/test_cloud_id.py b/cloudinit/cmd/tests/test_cloud_id.py
new file mode 100644
index 00000000..73738170
--- /dev/null
+++ b/cloudinit/cmd/tests/test_cloud_id.py
@@ -0,0 +1,127 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+"""Tests for cloud-id command line utility."""
+
+from cloudinit import util
+from collections import namedtuple
+from six import StringIO
+
+from cloudinit.cmd import cloud_id
+
+from cloudinit.tests.helpers import CiTestCase, mock
+
+
+class TestCloudId(CiTestCase):
+
+    args = namedtuple('cloudidargs', ('instance_data json long'))
+
+    def setUp(self):
+        super(TestCloudId, self).setUp()
+        self.tmp = self.tmp_dir()
+        self.instance_data = self.tmp_path('instance-data.json', dir=self.tmp)
+
+    def test_cloud_id_arg_parser_defaults(self):
+        """Validate the argument defaults when not provided by the end-user."""
+        cmd = ['cloud-id']
+        with mock.patch('sys.argv', cmd):
+            args = cloud_id.get_parser().parse_args()
+        self.assertEqual(
+            '/run/cloud-init/instance-data.json',
+            args.instance_data)
+        self.assertEqual(False, args.long)
+        self.assertEqual(False, args.json)
+
+    def test_cloud_id_arg_parse_overrides(self):
+        """Override argument defaults by specifying values for each param."""
+        util.write_file(self.instance_data, '{}')
+        cmd = ['cloud-id', '--instance-data', self.instance_data, '--long',
+               '--json']
+        with mock.patch('sys.argv', cmd):
+            args = cloud_id.get_parser().parse_args()
+        self.assertEqual(self.instance_data, args.instance_data)
+        self.assertEqual(True, args.long)
+        self.assertEqual(True, args.json)
+
+    def test_cloud_id_missing_instance_data_json(self):
+        """Exit error when the provided instance-data.json does not exist."""
+        cmd = ['cloud-id', '--instance-data', self.instance_data]
+        with mock.patch('sys.argv', cmd):
+            with mock.patch('sys.stderr', new_callable=StringIO) as m_stderr:
+                with self.assertRaises(SystemExit) as context_manager:
+                    cloud_id.main()
+        self.assertEqual(1, context_manager.exception.code)
+        self.assertIn(
+            "ERROR: File not found '%s'" % self.instance_data,
+            m_stderr.getvalue())
+
+    def test_cloud_id_non_json_instance_data(self):
+        """Exit error when the provided instance-data.json is not json."""
+        cmd = ['cloud-id', '--instance-data', self.instance_data]
+        util.write_file(self.instance_data, '{')
+        with mock.patch('sys.argv', cmd):
+            with mock.patch('sys.stderr', new_callable=StringIO) as m_stderr:
+                with self.assertRaises(SystemExit) as context_manager:
+                    cloud_id.main()
+        self.assertEqual(1, context_manager.exception.code)
+        self.assertIn(
+            "ERROR: File '%s' is not valid json." % self.instance_data,
+            m_stderr.getvalue())
+
+    def test_cloud_id_from_cloud_name_in_instance_data(self):
+        """Report canonical cloud-id from cloud_name in instance-data."""
+        util.write_file(
+            self.instance_data,
+            '{"v1": {"cloud_name": "mycloud", "region": "somereg"}}')
+        cmd = ['cloud-id', '--instance-data', self.instance_data]
+        with mock.patch('sys.argv', cmd):
+            with mock.patch('sys.stdout', new_callable=StringIO) as m_stdout:
+                with self.assertRaises(SystemExit) as context_manager:
+                    cloud_id.main()
+        self.assertEqual(0, context_manager.exception.code)
+        self.assertEqual("mycloud\n", m_stdout.getvalue())
+
+    def test_cloud_id_long_name_from_instance_data(self):
+        """Report long cloud-id format from cloud_name and region."""
+        util.write_file(
+            self.instance_data,
+            '{"v1": {"cloud_name": "mycloud", "region": "somereg"}}')
+        cmd = ['cloud-id', '--instance-data', self.instance_data, '--long']
+        with mock.patch('sys.argv', cmd):
+            with mock.patch('sys.stdout', new_callable=StringIO) as m_stdout:
+                with self.assertRaises(SystemExit) as context_manager:
+                    cloud_id.main()
+        self.assertEqual(0, context_manager.exception.code)
+        self.assertEqual("mycloud\tsomereg\n", m_stdout.getvalue())
+
+    def test_cloud_id_lookup_from_instance_data_region(self):
+        """Report discovered canonical cloud_id when region lookup matches."""
+        util.write_file(
+            self.instance_data,
+            '{"v1": {"cloud_name": "aws", "region": "cn-north-1",'
+            ' "platform": "ec2"}}')
+        cmd = ['cloud-id', '--instance-data', self.instance_data, '--long']
+        with mock.patch('sys.argv', cmd):
+            with mock.patch('sys.stdout', new_callable=StringIO) as m_stdout:
+                with self.assertRaises(SystemExit) as context_manager:
+                    cloud_id.main()
+        self.assertEqual(0, context_manager.exception.code)
+        self.assertEqual("aws-china\tcn-north-1\n", m_stdout.getvalue())
+
+    def test_cloud_id_lookup_json_instance_data_adds_cloud_id_to_json(self):
+        """Report v1 instance-data content with cloud_id when --json set."""
+        util.write_file(
+            self.instance_data,
+            '{"v1": {"cloud_name": "unknown", "region": "dfw",'
+            ' "platform": "openstack", "public_ssh_keys": []}}')
+        expected = util.json_dumps({
+            'cloud_id': 'openstack', 'cloud_name': 'unknown',
+            'platform': 'openstack', 'public_ssh_keys': [], 'region': 'dfw'})
+        cmd = ['cloud-id', '--instance-data', self.instance_data, '--json']
+        with mock.patch('sys.argv', cmd):
+            with mock.patch('sys.stdout', new_callable=StringIO) as m_stdout:
+                with self.assertRaises(SystemExit) as context_manager:
+                    cloud_id.main()
+        self.assertEqual(0, context_manager.exception.code)
+        self.assertEqual(expected + '\n', m_stdout.getvalue())
+
+# vi: ts=4 expandtab
diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index 9b90680f..e6966b31 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -58,6 +58,14 @@ METADATA_UNKNOWN = 'unknown'
 
 LOG = logging.getLogger(__name__)
 
+# CLOUD_ID_REGION_PREFIX_MAP format is:
+#  <region-match-prefix>: (<new-cloud-id>: <test_allowed_cloud_callable>)
+CLOUD_ID_REGION_PREFIX_MAP = {
+    'cn-': ('aws-china', lambda c: c == 'aws'),    # only change aws regions
+    'us-gov-': ('aws-gov', lambda c: c == 'aws'),  # only change aws regions
+    'china': ('azure-china', lambda c: c == 'azure'),  # only change azure
+}
+
 
 class DataSourceNotFoundException(Exception):
     pass
@@ -770,6 +778,25 @@ def instance_id_matches_system_uuid(instance_id, field='system-uuid'):
     return instance_id.lower() == dmi_value.lower()
 
 
+def canonical_cloud_id(cloud_name, region, platform):
+    """Lookup the canonical cloud-id for a given cloud_name and region."""
+    if not cloud_name:
+        cloud_name = METADATA_UNKNOWN
+    if not region:
+        region = METADATA_UNKNOWN
+    if region == METADATA_UNKNOWN:
+        if cloud_name != METADATA_UNKNOWN:
+            return cloud_name
+        return platform
+    for prefix, cloud_id_test in CLOUD_ID_REGION_PREFIX_MAP.items():
+        (cloud_id, valid_cloud) = cloud_id_test
+        if region.startswith(prefix) and valid_cloud(cloud_name):
+            return cloud_id
+    if cloud_name != METADATA_UNKNOWN:
+        return cloud_name
+    return platform
+
+
 def convert_vendordata(data, recurse=True):
     """data: a loaded object (strings, arrays, dicts).
     return something suitable for cloudinit vendordata_raw.
diff --git a/cloudinit/sources/tests/test_init.py b/cloudinit/sources/tests/test_init.py
index 391b3436..6378e98b 100644
--- a/cloudinit/sources/tests/test_init.py
+++ b/cloudinit/sources/tests/test_init.py
@@ -11,7 +11,8 @@ from cloudinit.helpers import Paths
 from cloudinit import importer
 from cloudinit.sources import (
     EXPERIMENTAL_TEXT, INSTANCE_JSON_FILE, INSTANCE_JSON_SENSITIVE_FILE,
-    REDACT_SENSITIVE_VALUE, UNSET, DataSource, redact_sensitive_keys)
+    METADATA_UNKNOWN, REDACT_SENSITIVE_VALUE, UNSET, DataSource,
+    canonical_cloud_id, redact_sensitive_keys)
 from cloudinit.tests.helpers import CiTestCase, skipIf, mock
 from cloudinit.user_data import UserDataProcessor
 from cloudinit import util
@@ -607,4 +608,75 @@ class TestRedactSensitiveData(CiTestCase):
             redact_sensitive_keys(md))
 
 
+class TestCanonicalCloudID(CiTestCase):
+
+    def test_cloud_id_returns_platform_on_unknowns(self):
+        """When region and cloud_name are unknown, return platform."""
+        self.assertEqual(
+            'platform',
+            canonical_cloud_id(cloud_name=METADATA_UNKNOWN,
+                               region=METADATA_UNKNOWN,
+                               platform='platform'))
+
+    def test_cloud_id_returns_platform_on_none(self):
+        """When region and cloud_name are unknown, return platform."""
+        self.assertEqual(
+            'platform',
+            canonical_cloud_id(cloud_name=None,
+                               region=None,
+                               platform='platform'))
+
+    def test_cloud_id_returns_cloud_name_on_unknown_region(self):
+        """When region is unknown, return cloud_name."""
+        for region in (None, METADATA_UNKNOWN):
+            self.assertEqual(
+                'cloudname',
+                canonical_cloud_id(cloud_name='cloudname',
+                                   region=region,
+                                   platform='platform'))
+
+    def test_cloud_id_returns_platform_on_unknown_cloud_name(self):
+        """When region is set but cloud_name is unknown return cloud_name."""
+        self.assertEqual(
+            'platform',
+            canonical_cloud_id(cloud_name=METADATA_UNKNOWN,
+                               region='region',
+                               platform='platform'))
+
+    def test_cloud_id_aws_based_on_region_and_cloud_name(self):
+        """When cloud_name is aws, return proper cloud-id based on region."""
+        self.assertEqual(
+            'aws-china',
+            canonical_cloud_id(cloud_name='aws',
+                               region='cn-north-1',
+                               platform='platform'))
+        self.assertEqual(
+            'aws',
+            canonical_cloud_id(cloud_name='aws',
+                               region='us-east-1',
+                               platform='platform'))
+        self.assertEqual(
+            'aws-gov',
+            canonical_cloud_id(cloud_name='aws',
+                               region='us-gov-1',
+                               platform='platform'))
+        self.assertEqual(  # Overrideen non-aws cloud_name is returned
+            '!aws',
+            canonical_cloud_id(cloud_name='!aws',
+                               region='us-gov-1',
+                               platform='platform'))
+
+    def test_cloud_id_azure_based_on_region_and_cloud_name(self):
+        """Report cloud-id when cloud_name is azure and region is in china."""
+        self.assertEqual(
+            'azure-china',
+            canonical_cloud_id(cloud_name='azure',
+                               region='chinaeast',
+                               platform='platform'))
+        self.assertEqual(
+            'azure',
+            canonical_cloud_id(cloud_name='azure',
+                               region='!chinaeast',
+                               platform='platform'))
+
 # vi: ts=4 expandtab
diff --git a/setup.py b/setup.py
index 5ed8eae2..ea37efc3 100755
--- a/setup.py
+++ b/setup.py
@@ -282,7 +282,8 @@ setuptools.setup(
     cmdclass=cmdclass,
     entry_points={
         'console_scripts': [
-            'cloud-init = cloudinit.cmd.main:main'
+            'cloud-init = cloudinit.cmd.main:main',
+            'cloud-id = cloudinit.cmd.cloud_id:main'
         ],
     }
 )