Retrieve SSH keys from IMDS first with OVF as a fallback (#509)

* pull ssh keys from imds first and fall back to ovf if unavailable

* refactor log and diagnostic messages

* refactor the OpenSSLManager instantiation and certificate usage

* fix unit test where exception was being silenced for generate cert

* fix tests now that certificate is not always generated

* add documentation for ssh key retrieval

* add ability to check if http client has security enabled

* refactor certificate logic to GoalState

2 files changed, 88 insertions, 15 deletions

diff --git a/cloudinit/sources/DataSourceAzure.py b/cloudinit/sources/DataSourceAzure.py
index f3c6452b..e98fd497 100755
--- a/cloudinit/sources/DataSourceAzure.py
+++ b/cloudinit/sources/DataSourceAzure.py
@@ -561,6 +561,40 @@ class DataSourceAzure(sources.DataSource):
     def device_name_to_device(self, name):
         return self.ds_cfg['disk_aliases'].get(name)
 
+    @azure_ds_telemetry_reporter
+    def get_public_ssh_keys(self):
+        """
+        Try to get the ssh keys from IMDS first, and if that fails
+        (i.e. IMDS is unavailable) then fallback to getting the ssh
+        keys from OVF.
+
+        The benefit to getting keys from IMDS is a large performance
+        advantage, so this is a strong preference. But we must keep
+        OVF as a second option for environments that don't have IMDS.
+        """
+        LOG.debug('Retrieving public SSH keys')
+        ssh_keys = []
+        try:
+            ssh_keys = [
+                public_key['keyData']
+                for public_key
+                in self.metadata['imds']['compute']['publicKeys']
+            ]
+            LOG.debug('Retrieved SSH keys from IMDS')
+        except KeyError:
+            log_msg = 'Unable to get keys from IMDS, falling back to OVF'
+            LOG.debug(log_msg)
+            report_diagnostic_event(log_msg)
+            try:
+                ssh_keys = self.metadata['public-keys']
+                LOG.debug('Retrieved keys from OVF')
+            except KeyError:
+                log_msg = 'No keys available from OVF'
+                LOG.debug(log_msg)
+                report_diagnostic_event(log_msg)
+
+        return ssh_keys
+
     def get_config_obj(self):
         return self.cfg
 
@@ -764,7 +798,22 @@ class DataSourceAzure(sources.DataSource):
         if self.ds_cfg['agent_command'] == AGENT_START_BUILTIN:
             self.bounce_network_with_azure_hostname()
 
-            pubkey_info = self.cfg.get('_pubkeys', None)
+            pubkey_info = None
+            try:
+                public_keys = self.metadata['imds']['compute']['publicKeys']
+                LOG.debug(
+                    'Successfully retrieved %s key(s) from IMDS',
+                    len(public_keys)
+                    if public_keys is not None
+                    else 0
+                )
+            except KeyError:
+                LOG.debug(
+                    'Unable to retrieve SSH keys from IMDS during '
+                    'negotiation, falling back to OVF'
+                )
+                pubkey_info = self.cfg.get('_pubkeys', None)
+
             metadata_func = partial(get_metadata_from_fabric,
                                     fallback_lease_file=self.
                                     dhclient_lease_file,
@@ -1443,7 +1492,7 @@ def get_metadata_from_imds(fallback_nic, retries):
 @azure_ds_telemetry_reporter
 def _get_metadata_from_imds(retries):
 
-    url = IMDS_URL + "instance?api-version=2017-12-01"
+    url = IMDS_URL + "instance?api-version=2019-06-01"
     headers = {"Metadata": "true"}
     try:
         response = readurl(
diff --git a/cloudinit/sources/helpers/azure.py b/cloudinit/sources/helpers/azure.py
index 507f6ac8..79445a81 100755
--- a/cloudinit/sources/helpers/azure.py
+++ b/cloudinit/sources/helpers/azure.py
@@ -288,12 +288,16 @@ class InvalidGoalStateXMLException(Exception):
 
 class GoalState:
 
-    def __init__(self, unparsed_xml: str,
-                 azure_endpoint_client: AzureEndpointHttpClient) -> None:
+    def __init__(
+            self,
+            unparsed_xml: str,
+            azure_endpoint_client: AzureEndpointHttpClient,
+            need_certificate: bool = True) -> None:
         """Parses a GoalState XML string and returns a GoalState object.
 
         @param unparsed_xml: string representing a GoalState XML.
-        @param azure_endpoint_client: instance of AzureEndpointHttpClient
+        @param azure_endpoint_client: instance of AzureEndpointHttpClient.
+        @param need_certificate: switch to know if certificates is needed.
         @return: GoalState object representing the GoalState XML string.
         """
         self.azure_endpoint_client = azure_endpoint_client
@@ -322,7 +326,7 @@ class GoalState:
         url = self._text_from_xpath(
             './Container/RoleInstanceList/RoleInstance'
             '/Configuration/Certificates')
-        if url is not None:
+        if url is not None and need_certificate:
             with events.ReportEventStack(
                     name="get-certificates-xml",
                     description="get certificates xml",
@@ -741,27 +745,38 @@ class WALinuxAgentShim:
             GoalState.
         @return: The list of user's authorized pubkey values.
         """
-        if self.openssl_manager is None:
+        http_client_certificate = None
+        if self.openssl_manager is None and pubkey_info is not None:
             self.openssl_manager = OpenSSLManager()
+            http_client_certificate = self.openssl_manager.certificate
         if self.azure_endpoint_client is None:
             self.azure_endpoint_client = AzureEndpointHttpClient(
-                self.openssl_manager.certificate)
-        goal_state = self._fetch_goal_state_from_azure()
-        ssh_keys = self._get_user_pubkeys(goal_state, pubkey_info)
+                http_client_certificate)
+        goal_state = self._fetch_goal_state_from_azure(
+            need_certificate=http_client_certificate is not None
+        )
+        ssh_keys = None
+        if pubkey_info is not None:
+            ssh_keys = self._get_user_pubkeys(goal_state, pubkey_info)
         health_reporter = GoalStateHealthReporter(
             goal_state, self.azure_endpoint_client, self.endpoint)
         health_reporter.send_ready_signal()
         return {'public-keys': ssh_keys}
 
     @azure_ds_telemetry_reporter
-    def _fetch_goal_state_from_azure(self) -> GoalState:
+    def _fetch_goal_state_from_azure(
+            self,
+            need_certificate: bool) -> GoalState:
         """Fetches the GoalState XML from the Azure endpoint, parses the XML,
         and returns a GoalState object.
 
         @return: GoalState object representing the GoalState XML
         """
         unparsed_goal_state_xml = self._get_raw_goal_state_xml_from_azure()
-        return self._parse_raw_goal_state_xml(unparsed_goal_state_xml)
+        return self._parse_raw_goal_state_xml(
+            unparsed_goal_state_xml,
+            need_certificate
+        )
 
     @azure_ds_telemetry_reporter
     def _get_raw_goal_state_xml_from_azure(self) -> str:
@@ -774,7 +789,11 @@ class WALinuxAgentShim:
         LOG.info('Registering with Azure...')
         url = 'http://{}/machine/?comp=goalstate'.format(self.endpoint)
         try:
-            response = self.azure_endpoint_client.get(url)
+            with events.ReportEventStack(
+                    name="goalstate-retrieval",
+                    description="retrieve goalstate",
+                    parent=azure_ds_reporter):
+                response = self.azure_endpoint_client.get(url)
         except Exception as e:
             msg = 'failed to register with Azure: %s' % e
             LOG.warning(msg)
@@ -785,7 +804,9 @@ class WALinuxAgentShim:
 
     @azure_ds_telemetry_reporter
     def _parse_raw_goal_state_xml(
-            self, unparsed_goal_state_xml: str) -> GoalState:
+            self,
+            unparsed_goal_state_xml: str,
+            need_certificate: bool) -> GoalState:
         """Parses a GoalState XML string and returns a GoalState object.
 
         @param unparsed_goal_state_xml: GoalState XML string
@@ -793,7 +814,10 @@ class WALinuxAgentShim:
         """
         try:
             goal_state = GoalState(
-                unparsed_goal_state_xml, self.azure_endpoint_client)
+                unparsed_goal_state_xml,
+                self.azure_endpoint_client,
+                need_certificate
+            )
         except Exception as e:
             msg = 'Error processing GoalState XML: %s' % e
             LOG.warning(msg)