Skip to content

Commit

Permalink
Turn back into binary agent or userdata if dict
Browse files Browse the repository at this point in the history
To properly assemble a multipart data we need raw bytes
If both are bytes just pass them as they are.
If either user or agent data is dict we need to assemble back
the raw binary data (including the content type stenza)

(Maybe we should reload the files in their binary form instead?)
  • Loading branch information
CarlosNihelton committed Jul 23, 2024
1 parent fe1629d commit 5019825
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 4 deletions.
19 changes: 15 additions & 4 deletions cloudinit/sources/DataSourceWSL.py
Original file line number Diff line number Diff line change
Expand Up @@ -320,13 +320,24 @@ def _get_data(self) -> bool:
if not any([user_data, agent_data]):
return False

# If we cannot reliably model data files as dicts, then we cannot merge
# ourselves, so we can pass the data in ascending order as a list for
# cloud-init to handle internally
if isinstance(agent_data, bytes) or isinstance(user_data, bytes):
# If we cannot reliably model both data files as dicts, then we cannot
# merge them ourselves, so we can pass the data in ascending order as
# a list of binary contents for cloud-init to handle internally
if isinstance(agent_data, bytes) and isinstance(user_data, bytes):
self.userdata_raw = cast(Any, [user_data, agent_data])
return True

if isinstance(user_data, bytes):
self.userdata_raw = cast(
Any, [user_data, "#cloud-config\n%s" % yaml.dump(agent_data)]
)
return True

if isinstance(agent_data, bytes):
self.userdata_raw = cast(
Any, ["#cloud-config\n%s" % yaml.dump(user_data), agent_data]
)
return True
# We only care about overriding modules entirely, so we can just
# iterate over the top level keys and write over them if the agent
# provides them instead.
Expand Down
61 changes: 61 additions & 0 deletions tests/unittests/sources/test_wsl.py
Original file line number Diff line number Diff line change
Expand Up @@ -625,6 +625,67 @@ def test_landscape_empty_data(self, m_get_linux_dist, tmpdir, paths):
"agent_test" in userdata and "agent_token" in userdata
), "Agent data should be present"

@mock.patch("cloudinit.util.get_linux_distro")
def test_landscape_shell_script(self, m_get_linux_dist, tmpdir, paths):
"""Asserts that Pro for WSL and Landscape goes multipart"""

m_get_linux_dist.return_value = SAMPLE_LINUX_DISTRO

ubuntu_pro_tmp = tmpdir.join(".ubuntupro", ".cloud-init")
os.makedirs(ubuntu_pro_tmp, exist_ok=True)

agent_file = ubuntu_pro_tmp.join("agent.yaml")
agent_file.write(
"""#cloud-config
landscape:
host:
url: hosted.com:6554
client:
account_name: agent_test
url: https://hosted.com/message-system
ping_url: https://hosted.com/ping
ssl_public_key: C:\\Users\\User\\server.pem
tags: wsl
ubuntu_pro:
token: agent_token"""
)

COMMAND = "echo Hello cloud-init on WSL!"
landscape_file = ubuntu_pro_tmp.join("%s.user-data" % INSTANCE_NAME)
landscape_file.write(f"#!/bin/sh\n{COMMAND}\n")

# Run the datasource
ds = wsl.DataSourceWSL(
sys_cfg=SAMPLE_CFG,
distro=_get_distro("ubuntu"),
paths=paths,
)

# Assert Landscape and Agent combine, with Agent taking precedence
assert ds.get_data() is True
ud = ds.get_userdata()

assert ud is not None
userdata = cast(
str,
join_payloads_from_content_type(
cast(MIMEMultipart, ud), "text/cloud-config"
),
)

assert (
"agent_test" in userdata and "agent_token" in userdata
), "Agent data should be present"

shell_script = cast(
str,
join_payloads_from_content_type(
cast(MIMEMultipart, ud), "text/x-shellscript"
),
)

assert COMMAND in shell_script

@mock.patch("cloudinit.util.get_linux_distro")
def test_with_landscape_no_tags(self, m_get_linux_dist, tmpdir, paths):
"""Validates the Pro For WSL default Landscape tags are applied"""
Expand Down

0 comments on commit 5019825

Please sign in to comment.