Skip to content

Implement fallback to WireServer when deploying SSH keys #3603

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
narrieta wants to merge 4 commits into
base: develop
Choose a base branch
from
+158 −19
Open

Implement fallback to WireServer when deploying SSH keys #3603

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a101540
Implement fallback to WireServer when deploying SSH keys
Apr 29, 2026
54d581f
comment
Apr 29, 2026
d1dd817
review
Apr 29, 2026
42f6b6f
Merge branch 'develop' into certificates
narrieta Apr 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions azurelinuxagent/pa/provision/default.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
from azurelinuxagent.common.exception import ProvisionError, ProtocolError, \
OSUtilError
from azurelinuxagent.common.osutil import get_osutil
from azurelinuxagent.common.protocol.goal_state import GoalState, GoalStateProperties
from azurelinuxagent.common.protocol.restapi import ProvisionStatus
from azurelinuxagent.common.protocol.util import get_protocol_util
from azurelinuxagent.common.version import AGENT_NAME
Expand Down Expand Up @@ -240,9 +241,34 @@ def config_user_account(self, ovfenv):
logger.info("Configure sshd")
self.osutil.conf_sshd(ovfenv.disable_ssh_password_auth)

self._download_ssh_keys(ovfenv)
self.deploy_ssh_pubkeys(ovfenv)
self.deploy_ssh_keypairs(ovfenv)

def _download_ssh_keys(self, ovfenv):
Copy link
Copy Markdown
Contributor

@nagworld9 nagworld9 May 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since its not download always, can we rename to _download_ssh_keys_if_needed or something like that?

#
Copy link
Copy Markdown
Member Author

@narrieta narrieta Apr 29, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the reason of this check, see the implementation of self.deploy_ssh_pubkeys() and self.deploy_ssh_keypairs(). In summary, the former will look for /var/lib/waagent/*.crt when ovfenv.xml doesn't have a value for the key, and the latter will look for /var/lib/waagent/*.prv when ovfenv.xml contains any key pairs.

For the format of ovfenv.xml, see the unit tests.

# We need to download the Certificates package from the Wireserver if any public key in ovfenv.xml has only a thumbprint (i.e. no value for the key) or if any key pairs need to be installed
#
download_certificates = any(value is None and thumbprint is not None for _, thumbprint, value in ovfenv.ssh_pubkeys) or len(ovfenv.ssh_keypairs) > 0
if not download_certificates:
return

#
Copy link
Copy Markdown
Member Author

@narrieta narrieta Apr 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Originally, the PA would retrieve the goal state as part of protocol detection, which retries 360 times, with a 10 second period. If that fails, the process exits and is restarted by the system, which goes again into the 360 retries, etc.

I feel just a few tries are enough, let me know if you think otherwise; thanks.

# Try to download the certificates a few times but continue execution if all attempts fail. The code that deploys the SSH keys will report that as a provisioning error.
#
max_attempts = 5
for attempt in range(max_attempts):
try:
protocol = self.protocol_util.get_protocol(init_goal_state=False)
_ = GoalState(protocol.client, goal_state_properties=GoalStateProperties.Certificates)
return
Comment thread
narrieta marked this conversation as resolved.
except Exception as e:
if attempt < max_attempts - 1:
logger.warn("Unable to download certificates; will retry after a short delay: {0}", ustr(e))
time.sleep(30)
else:
Comment thread
narrieta marked this conversation as resolved.
logger.error("Unable to download certificates: {0}", ustr(e))

def save_customdata(self, ovfenv):
customdata = ovfenv.customdata
if customdata is None:
Expand Down
6 changes: 0 additions & 6 deletions tests/data/ovf-env-2.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,6 @@
<Value>ssh-rsa AAAANOTAREALKEY== foo@bar.local</Value>
</PublicKey>
</PublicKeys>
<KeyPairs>
Copy link
Copy Markdown
Member Author

@narrieta narrieta Apr 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The key pairs in the unit tests are not used, and with the new code they would trigger the Certificates download. To fix this, quite a few tests need to be rewritten to mock that download. Since the keys are not used by the tests, I decided to remove them instead.

<KeyPair>
<Fingerprint>EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62</Fingerprint>
<Path>$HOME/UserName/.ssh/id_rsa</Path>
</KeyPair>
</KeyPairs>
</SSH>
<CustomData>CustomData</CustomData>
</LinuxProvisioningConfigurationSet>
Expand Down
6 changes: 0 additions & 6 deletions tests/data/ovf-env-4.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,6 @@
<Value>ssh-rsa AAAANOTAREALKEY== foo@bar.local</Value>
</PublicKey>
</PublicKeys>
<KeyPairs>
<KeyPair>
<Fingerprint>EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62</Fingerprint>
<Path>$HOME/UserName/.ssh/id_rsa</Path>
</KeyPair>
</KeyPairs>
</SSH>
<CustomData>CustomData</CustomData>
</LinuxProvisioningConfigurationSet>
Expand Down
6 changes: 0 additions & 6 deletions tests/data/ovf-env.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,6 @@
<Value>ssh-rsa AAAANOTAREALKEY== foo@bar.local</Value>
</PublicKey>
</PublicKeys>
<KeyPairs>
<KeyPair>
<Fingerprint>EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62</Fingerprint>
<Path>$HOME/UserName/.ssh/id_rsa</Path>
</KeyPair>
</KeyPairs>
</SSH>
<CustomData>CustomData</CustomData>
</LinuxProvisioningConfigurationSet>
Expand Down
33 changes: 33 additions & 0 deletions tests/data/ovf-env_key_pair.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
<?xml version="1.0" encoding="utf-8"?>
<Environment xmlns="http://schemas.dmtf.org/ovf/environment/1" xmlns:oe="http://schemas.dmtf.org/ovf/environment/1" xmlns:wa="http://schemas.microsoft.com/windowsazure" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
<wa:ProvisioningSection>
<wa:Version>1.0</wa:Version>
<LinuxProvisioningConfigurationSet xmlns="http://schemas.microsoft.com/windowsazure" xmlns:i="http://www.w3.org/2001/XMLSchema-instance">
<ConfigurationSetType>LinuxProvisioningConfiguration</ConfigurationSetType>
<HostName>HostName</HostName>
<UserName>UserName</UserName>
<UserPassword>UserPassword</UserPassword>
<DisableSshPasswordAuthentication>false</DisableSshPasswordAuthentication>
<SSH>
<KeyPairs>
<KeyPair>
<Fingerprint>EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62</Fingerprint>
<Path>$HOME/UserName/.ssh/id_rsa</Path>
</KeyPair>
</KeyPairs>
</SSH>
<CustomData>CustomData</CustomData>
</LinuxProvisioningConfigurationSet>
</wa:ProvisioningSection>
<wa:PlatformSettingsSection>
<wa:Version>1.0</wa:Version>
<wa:PlatformSettings>
<wa:KmsServerHostname>kms.core.windows.net</wa:KmsServerHostname>
<wa:ProvisionGuestAgent>false</wa:ProvisionGuestAgent>
<wa:GuestAgentPackageName xsi:nil="true"/>
<wa:RetainWindowsPEPassInUnattend>true</wa:RetainWindowsPEPassInUnattend>
<wa:RetainOfflineServicingPassInUnattend>true</wa:RetainOfflineServicingPassInUnattend>
<wa:PreprovisionedVm>false</wa:PreprovisionedVm>
</wa:PlatformSettings>
</wa:PlatformSettingsSection>
</Environment>
34 changes: 34 additions & 0 deletions tests/data/ovf-env_public_key.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
<?xml version="1.0" encoding="utf-8"?>
<Environment xmlns="http://schemas.dmtf.org/ovf/environment/1" xmlns:oe="http://schemas.dmtf.org/ovf/environment/1" xmlns:wa="http://schemas.microsoft.com/windowsazure" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
<wa:ProvisioningSection>
<wa:Version>1.0</wa:Version>
<LinuxProvisioningConfigurationSet xmlns="http://schemas.microsoft.com/windowsazure" xmlns:i="http://www.w3.org/2001/XMLSchema-instance">
<ConfigurationSetType>LinuxProvisioningConfiguration</ConfigurationSetType>
<HostName>HostName</HostName>
<UserName>UserName</UserName>
<UserPassword>UserPassword</UserPassword>
<DisableSshPasswordAuthentication>false</DisableSshPasswordAuthentication>
<SSH>
<PublicKeys>
<PublicKey>
<Fingerprint>EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62</Fingerprint>
Copy link
Copy Markdown
Contributor

@maddieford maddieford May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it expected that the customer should only provide one of fingerprint and value, or is it acceptable to provide both?

Copy link
Copy Markdown
Member Author

@narrieta narrieta May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The original sample file from which I copied this one includes both. In practice, I believe only 1 would be populated.

Now, the logic in OSUtil.deploy_ssh_pubkey() prioritizes the value over the thumbprint, so in terms of the unit test that uses this new file, having both is a good test case.

<Path>$HOME/UserName/.ssh/authorized_keys</Path>
<Value>ssh-rsa AAAANOTAREALKEY== foo@bar.local</Value>
</PublicKey>
</PublicKeys>
</SSH>
<CustomData>CustomData</CustomData>
</LinuxProvisioningConfigurationSet>
</wa:ProvisioningSection>
<wa:PlatformSettingsSection>
<wa:Version>1.0</wa:Version>
<wa:PlatformSettings>
<wa:KmsServerHostname>kms.core.windows.net</wa:KmsServerHostname>
<wa:ProvisionGuestAgent>false</wa:ProvisionGuestAgent>
<wa:GuestAgentPackageName xsi:nil="true"/>
<wa:RetainWindowsPEPassInUnattend>true</wa:RetainWindowsPEPassInUnattend>
<wa:RetainOfflineServicingPassInUnattend>true</wa:RetainOfflineServicingPassInUnattend>
<wa:PreprovisionedVm>false</wa:PreprovisionedVm>
</wa:PlatformSettings>
</wa:PlatformSettingsSection>
</Environment>
33 changes: 33 additions & 0 deletions tests/data/ovf-env_public_key_no_value.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
<?xml version="1.0" encoding="utf-8"?>
<Environment xmlns="http://schemas.dmtf.org/ovf/environment/1" xmlns:oe="http://schemas.dmtf.org/ovf/environment/1" xmlns:wa="http://schemas.microsoft.com/windowsazure" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
<wa:ProvisioningSection>
<wa:Version>1.0</wa:Version>
<LinuxProvisioningConfigurationSet xmlns="http://schemas.microsoft.com/windowsazure" xmlns:i="http://www.w3.org/2001/XMLSchema-instance">
<ConfigurationSetType>LinuxProvisioningConfiguration</ConfigurationSetType>
<HostName>HostName</HostName>
<UserName>UserName</UserName>
<UserPassword>UserPassword</UserPassword>
<DisableSshPasswordAuthentication>false</DisableSshPasswordAuthentication>
<SSH>
<PublicKeys>
<PublicKey>
<Fingerprint>EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62</Fingerprint>
<Path>$HOME/UserName/.ssh/authorized_keys</Path>
</PublicKey>
</PublicKeys>
</SSH>
<CustomData>CustomData</CustomData>
</LinuxProvisioningConfigurationSet>
</wa:ProvisioningSection>
<wa:PlatformSettingsSection>
<wa:Version>1.0</wa:Version>
<wa:PlatformSettings>
<wa:KmsServerHostname>kms.core.windows.net</wa:KmsServerHostname>
<wa:ProvisionGuestAgent>false</wa:ProvisionGuestAgent>
<wa:GuestAgentPackageName xsi:nil="true"/>
<wa:RetainWindowsPEPassInUnattend>true</wa:RetainWindowsPEPassInUnattend>
<wa:RetainOfflineServicingPassInUnattend>true</wa:RetainOfflineServicingPassInUnattend>
<wa:PreprovisionedVm>false</wa:PreprovisionedVm>
</wa:PlatformSettings>
</wa:PlatformSettingsSection>
</Environment>
33 changes: 32 additions & 1 deletion tests/pa/test_provision.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,20 +14,23 @@
#
# Requires Python 2.6+ and Openssl 1.0+
#

import contextlib
import os
import re
import unittest

import azurelinuxagent.common.conf as conf
from azurelinuxagent.common.exception import ProvisionError
from azurelinuxagent.common.osutil.default import DefaultOSUtil
from azurelinuxagent.common.protocol.ovfenv import OvfEnv
from azurelinuxagent.common.protocol.util import OVF_FILE_NAME
from azurelinuxagent.pa.provision import get_provision_handler
from azurelinuxagent.pa.provision.cloudinit import CloudInitProvisionHandler
from azurelinuxagent.pa.provision.default import ProvisionHandler
from azurelinuxagent.common.utils import fileutil
from tests.lib import wire_protocol_data
from tests.lib.tools import AgentTestCase, distros, load_data, MagicMock, Mock, patch
from tests.lib.mock_wire_protocol import mock_wire_protocol


class TestProvision(AgentTestCase):
Expand Down Expand Up @@ -377,6 +380,34 @@ def test_get_provision_handler_config_cloudinit(
self.assertIsInstance(provisioning_handler, CloudInitProvisionHandler, 'Provisioning handler should be cloud-init if agent is set to cloud-init')


@staticmethod
@contextlib.contextmanager
def _create_provision_handler_with_mock_protocol():
handler = ProvisionHandler()
handler.protocol_util = Mock()

with mock_wire_protocol(wire_protocol_data.DATA_FILE, detect_protocol=False) as mock_protocol:
handler.protocol_util.get_protocol = Mock(return_value=mock_protocol)
yield handler, mock_protocol

def test_it_should_not_download_certificates_when_the_public_key_has_a_value(self):
ovfenv = OvfEnv(load_data("ovf-env_public_key.xml"))
with TestProvision._create_provision_handler_with_mock_protocol() as (handler, protocol):
handler._download_ssh_keys(ovfenv)
self.assertEqual(0, protocol.mock_wire_data.call_counts['certificates'], "The Certificates package should not have been retrieved")

def test_it_should_download_certificates_when_the_public_key_does_not_have_a_value(self):
ovfenv = OvfEnv(load_data("ovf-env_public_key_no_value.xml"))
with TestProvision._create_provision_handler_with_mock_protocol() as (handler, protocol):
handler._download_ssh_keys(ovfenv)
self.assertEqual(1, protocol.mock_wire_data.call_counts['certificates'], "The Certificates package should have been retrieved")

def test_it_should_download_certificates_when_key_pairs_need_to_be_deployed(self):
ovfenv = OvfEnv(load_data("ovf-env_key_pair.xml"))
Comment thread
narrieta marked this conversation as resolved.
with TestProvision._create_provision_handler_with_mock_protocol() as (handler, protocol):
handler._download_ssh_keys(ovfenv)
self.assertEqual(1, protocol.mock_wire_data.call_counts['certificates'], "The Certificates package should have been retrieved")

Copy link
Copy Markdown
Contributor

@nagworld9 nagworld9 May 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cam we add test around exception when downloading certs?

Copy link
Copy Markdown
Member Author

@narrieta narrieta May 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Those are covered in the unit tests for the goal state.

Are you thinking of a specific scenario?

Copy link
Copy Markdown
Member Author

@narrieta narrieta May 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Or do you mean the retry logic in _download_ssh_keys?

Copy link
Copy Markdown
Contributor

@nagworld9 nagworld9 May 4, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Retry logic in download_ssh_keys. We expect next operation would raise error if we failed to download with all retries.

can we test that we don't re-raise error here? It would be good to test if next operation rising that provision error if not already present in the unit tests

Copy link
Copy Markdown
Member Author

@narrieta narrieta May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, i can add a test for the retry logic

if __name__ == '__main__':
unittest.main()

Loading