j3d1/webvirtcloud
1
0
Fork 0
mirror of https://github.com/retspen/webvirtcloud synced 2026-11-08 11:16:47 +00:00
Code Issues Projects Releases Wiki Activity

feat(security): harden storage uploads and vrtManager snapshot XML handling

Browse source
This commit is contained in:
catborise 2026-09-27 13:28:17 +03:00
parent b3e6fb2b7a
commit 5e240236c3
5 changed files with 389 additions and 104 deletions
Download patch file Download diff file Expand all files Collapse all files

119
storages/tests.py
View file

@ -1,6 +1,121 @@
import os
import shutil
import tempfile
from unittest.mock import MagicMock, patch
from django.core.files.uploadedfile import SimpleUploadedFile
from django.test import TestCase
from django.urls import reverse
from django.contrib.auth import get_user_model
from computes.models import Compute
from storages.views import handle_uploaded_file
from vrtManager.connection import CONN_SOCKET, CONN_SSH
User = get_user_model()
class StoragesTestCase(TestCase):
class StorageUploadSecurityTests(TestCase):
def setUp(self):
pass
self.temp_dir = tempfile.mkdtemp()
self.mock_socket_conn = MagicMock()
self.mock_socket_conn.conn = CONN_SOCKET
self.mock_ssh_conn = MagicMock()
self.mock_ssh_conn.conn = CONN_SSH
self.mock_ssh_conn.host = "127.0.0.1"
self.mock_ssh_conn.login = "test"
self.mock_ssh_conn.passwd = "test"
def tearDown(self):
shutil.rmtree(self.temp_dir, ignore_errors=True)
def test_handle_uploaded_file_invalid_names(self):
dummy_chunk = SimpleUploadedFile("test.iso", b"data")
invalid_names = ["", " ", ".", "..", "test\x00evil.iso", "../evil.iso", "../../etc/shadow"]
for bad_name in invalid_names:
clean_base = os.path.basename(bad_name).strip()
# bad_name either gets rejected as empty/dot/null or has basename extracted
if not clean_base or clean_base in (".", "..") or "\x00" in clean_base:
with self.assertRaises(ValueError):
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, bad_name, dummy_chunk, is_last_chunk=True)
def test_handle_uploaded_file_successful_socket_upload(self):
dummy_chunk = SimpleUploadedFile("install.iso", b"ISODATA_BLOCK_123")
file_name = "install.iso"
# Chunk 1 (is_last_chunk=False)
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, dummy_chunk, is_last_chunk=False)
part_path = os.path.join(self.temp_dir, f"{file_name}.part")
final_path = os.path.join(self.temp_dir, file_name)
self.assertTrue(os.path.exists(part_path))
self.assertFalse(os.path.exists(final_path))
# Chunk 2 (is_last_chunk=True)
dummy_chunk2 = SimpleUploadedFile("install.iso", b"_PART2")
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, dummy_chunk2, is_last_chunk=True)
self.assertFalse(os.path.exists(part_path))
self.assertTrue(os.path.exists(final_path))
with open(final_path, "rb") as f:
self.assertEqual(f.read(), b"ISODATA_BLOCK_123_PART2")
def test_handle_uploaded_file_jail_containment(self):
dummy_chunk = SimpleUploadedFile("test.iso", b"data")
# Attempt traversal outside sandbox
with patch("os.path.basename", side_effect=lambda x: x): # simulate raw traversal name
with self.assertRaises(PermissionError):
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, "../../../evil.iso", dummy_chunk, is_last_chunk=True)
def test_handle_uploaded_file_ssh_traversal_check(self):
dummy_chunk = SimpleUploadedFile("test.iso", b"data")
with patch("os.path.basename", side_effect=lambda x: x):
with self.assertRaises(PermissionError):
handle_uploaded_file(self.mock_ssh_conn, "/var/lib/libvirt/images", "../../../evil.iso", dummy_chunk, is_last_chunk=True)
class StorageUploadViewTests(TestCase):
def setUp(self):
self.admin = User.objects.create_superuser(username="storage_admin", password="password", email="storage_admin@example.com")
self.client.login(username="storage_admin", password="password")
self.compute = Compute.objects.create(name="LocalCompute", hostname="localhost", type=1, login="u", password="p")
@patch("storages.views.wvmStorage")
def test_iso_upload_missing_chunk(self, mock_wvm):
mock_conn = MagicMock()
mock_conn.get_status.return_value = 1
mock_conn.get_target_path.return_value = "/tmp"
mock_conn.get_type.return_value = "dir"
mock_conn.get_autostart.return_value = 1
mock_conn.get_size.return_value = [1000, 500]
mock_conn.update_volumes.return_value = []
mock_wvm.return_value = mock_conn
url = reverse("storage", args=[self.compute.id, "default"])
res = self.client.post(url, {"iso_upload": "1", "file_name": "test.iso"})
self.assertEqual(res.status_code, 400)
self.assertIn("No file chunk was submitted", res.json().get("error", ""))
@patch("storages.views.wvmStorage")
def test_iso_upload_invalid_filename(self, mock_wvm):
mock_conn = MagicMock()
mock_conn.get_status.return_value = 1
mock_conn.get_target_path.return_value = "/tmp"
mock_conn.get_type.return_value = "dir"
mock_conn.get_autostart.return_value = 1
mock_conn.get_size.return_value = [1000, 500]
mock_conn.update_volumes.return_value = []
mock_wvm.return_value = mock_conn
url = reverse("storage", args=[self.compute.id, "default"])
chunk = SimpleUploadedFile("test.iso", b"chunkdata")
res = self.client.post(url, {
"iso_upload": "1",
"file": chunk,
"file_name": "..\x00..",
"chunk_index": "0",
"total_chunks": "1",
})
self.assertEqual(res.status_code, 400)
self.assertIn("Invalid file name", res.json().get("error", ""))

172
storages/views.py
View file

@ -1,5 +1,6 @@
import json
import os
import posixpath
from admin.decorators import superuser_only
from appsettings.settings import app_settings
@ -97,6 +98,74 @@ def storages(request, compute_id):
return render(request, "storages.html", locals())
def handle_uploaded_file(conn, path, file_name, file_chunk, is_last_chunk):
clean_name = os.path.basename(file_name).strip()
if not clean_name or clean_name in (".", "..") or "\x00" in clean_name:
raise ValueError(_("Invalid file name"))
if conn.conn == CONN_SSH:
remote_base = posixpath.normpath(path)
target_temp = posixpath.normpath(posixpath.join(remote_base, f"{clean_name}.part"))
target_final = posixpath.normpath(posixpath.join(remote_base, clean_name))
if not (target_temp.startswith(remote_base + "/") or target_temp == remote_base) or \
not (target_final.startswith(remote_base + "/") or target_final == remote_base):
raise PermissionError(_("Security issues with file uploading: path traversal detected"))
try:
hostname, port = conn.host, 22
if ":" in hostname:
hostname, port_str = hostname.split(":")
port = int(port_str)
ssh = paramiko.SSHClient()
ssh.load_system_host_keys()
ssh.set_missing_host_key_policy(paramiko.WarningPolicy())
ssh.connect(hostname=hostname, port=port, username=conn.login, password=conn.passwd)
try:
sftp = ssh.open_sftp()
try:
remote_file = sftp.open(target_temp, "ab")
try:
remote_file.set_pipelined(True)
for chunk_data in file_chunk.chunks():
remote_file.write(chunk_data)
finally:
remote_file.close()
if is_last_chunk:
sftp.rename(target_temp, target_final)
finally:
sftp.close()
finally:
ssh.close()
except Exception as e:
raise Exception(_("SSH upload failed: {}").format(e))
elif conn.conn == CONN_SOCKET:
base_dir = os.path.abspath(path)
target_temp = os.path.abspath(os.path.join(base_dir, f"{clean_name}.part"))
target_final = os.path.abspath(os.path.join(base_dir, clean_name))
if (
os.path.commonpath([base_dir, target_temp]) != base_dir
or os.path.commonpath([base_dir, target_final]) != base_dir
):
raise PermissionError(_("Security issues with file uploading: path traversal detected"))
try:
with open(target_temp, "ab") as f:
for chunk_data in file_chunk.chunks():
f.write(chunk_data)
if is_last_chunk:
if os.path.exists(target_final):
os.remove(target_final)
os.rename(target_temp, target_final)
except FileNotFoundError:
raise Exception(_("File not found. Check the path variable and filename"))
else:
raise Exception(_("Unsupported connection type for file upload."))
@superuser_only
def storage(request, compute_id, pool):
"""
@ -105,53 +174,6 @@ def storage(request, compute_id, pool):
:param pool:
:return:
"""
def handle_uploaded_file(conn, path, file_name, file_chunk, is_last_chunk):
temp_name = f"{file_name}.part"
target_temp = os.path.normpath(os.path.join(path, temp_name))
target_final = os.path.normpath(os.path.join(path, file_name))
if not target_temp.startswith(path) or not target_final.startswith(path):
raise Exception(_("Security Issues with file uploading"))
if conn.conn == CONN_SSH:
try:
hostname, port = conn.host, 22
if ":" in hostname:
hostname, port_str = hostname.split(":")
port = int(port_str)
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(hostname=hostname, port=port, username=conn.login, password=conn.passwd)
sftp = ssh.open_sftp()
remote_file = sftp.open(target_temp, 'ab')
remote_file.set_pipelined(True)
for chunk_data in file_chunk.chunks():
remote_file.write(chunk_data)
remote_file.close()
if is_last_chunk:
sftp.rename(target_temp, target_final)
sftp.close()
ssh.close()
except Exception as e:
raise Exception(_("SSH upload failed: {}").format(e))
elif conn.conn == CONN_SOCKET:
try:
with open(target_temp, "ab") as f:
for chunk_data in file_chunk.chunks():
f.write(chunk_data)
if is_last_chunk:
if os.path.exists(target_final):
os.remove(target_final)
os.rename(target_temp, target_final)
except FileNotFoundError:
raise Exception(_("File not found. Check the path variable and filename"))
else:
raise Exception(_("Unsupported connection type for file upload."))
compute = get_object_or_404(Compute, pk=compute_id)
meta_prealloc = False
form = CreateVolumeForm()
@ -213,7 +235,12 @@ def storage(request, compute_id, pool):
if not file_chunk:
return JsonResponse({"error": _("No file chunk was submitted.")}, status=400)
file_name = request.POST.get("file_name")
file_name = request.POST.get("file_name", "")
clean_file_name = os.path.basename(file_name).strip()
if not clean_file_name or clean_file_name in (".", "..") or "\x00" in clean_file_name:
return JsonResponse({"error": _("Invalid file name.")}, status=400)
file_name = clean_file_name
chunk_index = int(request.POST.get("chunk_index", 0))
total_chunks = int(request.POST.get("total_chunks", 1))
is_last_chunk = chunk_index == total_chunks - 1
@ -223,28 +250,37 @@ def storage(request, compute_id, pool):
if file_name in conn.get_volumes():
return JsonResponse({"error": _("ISO image already exists")}, status=400)
# Clean up any partial files from previous failed uploads
temp_part_file = os.path.normpath(os.path.join(path, f"{file_name}.part"))
if conn.conn == CONN_SOCKET and os.path.exists(temp_part_file):
os.remove(temp_part_file)
if conn.conn == CONN_SOCKET:
base_dir = os.path.abspath(path)
temp_part_file = os.path.abspath(os.path.join(base_dir, f"{file_name}.part"))
if os.path.commonpath([base_dir, temp_part_file]) == base_dir and os.path.exists(temp_part_file):
os.remove(temp_part_file)
elif conn.conn == CONN_SSH:
try:
hostname, port = conn.host, 22
if ":" in hostname:
hostname, port_str = hostname.split(":")
port = int(port_str)
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(hostname=hostname, port=port, username=conn.login, password=conn.passwd)
sftp = ssh.open_sftp()
remote_base = posixpath.normpath(path)
temp_part_file = posixpath.normpath(posixpath.join(remote_base, f"{file_name}.part"))
if temp_part_file.startswith(remote_base + "/") or temp_part_file == remote_base:
try:
sftp.remove(temp_part_file)
except FileNotFoundError:
pass # File doesn't exist, which is fine
sftp.close()
ssh.close()
except Exception:
# Best effort to clean up, if it fails, let it be.
pass
hostname, port = conn.host, 22
if ":" in hostname:
hostname, port_str = hostname.split(":")
port = int(port_str)
ssh = paramiko.SSHClient()
ssh.load_system_host_keys()
ssh.set_missing_host_key_policy(paramiko.WarningPolicy())
ssh.connect(hostname=hostname, port=port, username=conn.login, password=conn.passwd)
try:
sftp = ssh.open_sftp()
try:
sftp.remove(temp_part_file)
except FileNotFoundError:
pass # File doesn't exist, which is fine
finally:
sftp.close()
finally:
ssh.close()
except Exception:
# Best effort to clean up, if it fails, let it be.
pass
try:
handle_uploaded_file(conn, path, file_name, file_chunk, is_last_chunk)

54
vrtManager/instance.py
View file

@ -1265,22 +1265,26 @@ class wvmInstance(wvmConnect):
def create_snapshot(self, name, desc=None):
state = "shutoff" if self.get_status() == 5 else "running"
name_esc = util.xml_escape(name)
desc_esc = util.xml_escape(desc) if desc else ""
xml = """<domainsnapshot>
<name>%s</name>
<description>%s</description>
<state>%s</state>
<creationTime>%d</creationTime>""" % (
name,
desc,
name_esc,
desc_esc,
state,
time.time(),
)
self.change_snapshot_xml()
xml += self._XMLDesc(VIR_DOMAIN_XML_SECURE)
xml += """<active>0</active>
</domainsnapshot>"""
self._snapshotCreateXML(xml, 0)
self.recover_snapshot_xml()
try:
xml += self._XMLDesc(VIR_DOMAIN_XML_SECURE)
xml += """<active>0</active>
</domainsnapshot>"""
self._snapshotCreateXML(xml, 0)
finally:
self.recover_snapshot_xml()
def change_snapshot_xml(self):
xml_temp = self._XMLDesc(VIR_DOMAIN_XML_SECURE).replace(
@ -1299,6 +1303,8 @@ class wvmInstance(wvmConnect):
def create_external_snapshot(self, name, date=None, desc=None):
creation_time = time.time()
state = "shutoff" if self.get_status() == 5 else "running"
name_esc = util.xml_escape(name)
desc_esc = util.xml_escape(desc) if desc else ""
#<seclabel type='none' model='dac' relabel='no'/>
xml = """<domainsnapshot>
<name>%s</name>
@ -1306,18 +1312,21 @@ class wvmInstance(wvmConnect):
<state>%s</state>
<creationTime>%d</creationTime>
""" % (
name,
desc,
name_esc,
desc_esc,
state,
creation_time,
)
self.change_snapshot_xml()
xml += self._XMLDesc(VIR_DOMAIN_XML_SECURE)
xml += """<active>0</active>
</domainsnapshot>"""
try:
xml += self._XMLDesc(VIR_DOMAIN_XML_SECURE)
xml += """<active>0</active>
</domainsnapshot>"""
self._snapshotCreateXML(xml, VIR_DOMAIN_SNAPSHOT_CREATE_DISK_ONLY)
self._snapshotCreateXML(xml, VIR_DOMAIN_SNAPSHOT_CREATE_DISK_ONLY)
finally:
self.recover_snapshot_xml()
self.refresh_instance_pools()
def get_external_snapshots(self):
@ -1394,9 +1403,11 @@ class wvmInstance(wvmConnect):
def snapshot_revert(self, snapshot):
self.change_snapshot_xml()
snap = self.instance.snapshotLookupByName(snapshot, 0)
self.instance.revertToSnapshot(snap, 0)
self.recover_snapshot_xml()
try:
snap = self.instance.snapshotLookupByName(snapshot, 0)
self.instance.revertToSnapshot(snap, 0)
finally:
self.recover_snapshot_xml()
def get_managed_save_image(self):
return self.instance.hasManagedSaveImage(0)
@ -1784,12 +1795,17 @@ class wvmInstance(wvmConnect):
return qos_values
def set_qos(self, mac, direction, average, peak, burst):
if direction not in ("inbound", "outbound"):
raise ValueError("Direction must be inbound or outbound")
average = int(average)
peak = int(peak)
burst = int(burst)
if direction == "inbound":
xml = f"<inbound average='{average}' peak='{peak}' burst='{burst}'/>"
elif direction == "outbound":
xml = f"<outbound average='{average}' peak='{peak}' burst='{burst}'/>"
else:
raise Exception("Direction must be inbound or outbound")
xml = f"<outbound average='{average}' peak='{peak}' burst='{burst}'/>"
tree = etree.fromstring(self._XMLDesc(0))

56
vrtManager/network.py
View file

@ -69,42 +69,55 @@ class wvmNetworks(wvmConnect):
openvswitch,
fixed=False,
):
esc_name = util.xml_escape(str(name))
esc_forward = util.xml_escape(str(forward)) if forward else ""
esc_bridge = util.xml_escape(str(bridge)) if bridge else ""
esc_gateway = util.xml_escape(str(gateway)) if gateway else ""
esc_mask = util.xml_escape(str(mask)) if mask else ""
esc_gateway6 = util.xml_escape(str(gateway6)) if gateway6 else ""
esc_prefix6 = util.xml_escape(str(prefix6)) if prefix6 else ""
xml = f"""
<network>
<name>{name}</name>"""
<name>{esc_name}</name>"""
if forward in ["nat", "route", "bridge"]:
xml += f"""<forward mode='{forward}'/>"""
xml += f"""<forward mode='{esc_forward}'/>"""
if forward == "macvtap":
xml += f"""<forward mode='bridge'>
<interface dev='{bridge}'/>
<interface dev='{esc_bridge}'/>
</forward>"""
else:
xml += """<bridge """
if forward in ["nat", "route", "none"]:
xml += """stp='on' delay='0'"""
if forward == "bridge":
xml += f"""name='{bridge}'"""
xml += f"""name='{esc_bridge}'"""
xml += """/>"""
if openvswitch is True:
xml += """<virtualport type='openvswitch'/>"""
if forward not in ["bridge", "macvtap"]:
if ipv4:
xml += f"""<ip address='{gateway}' netmask='{mask}'>"""
xml += f"""<ip address='{esc_gateway}' netmask='{esc_mask}'>"""
if dhcp4:
dhcp4_0 = util.xml_escape(str(dhcp4[0]))
dhcp4_1 = util.xml_escape(str(dhcp4[1]))
xml += f"""<dhcp>
<range start='{dhcp4[0]}' end='{dhcp4[1]}' />"""
<range start='{dhcp4_0}' end='{dhcp4_1}' />"""
if fixed:
fist_oct = int(dhcp4[0].strip().split(".")[3])
last_oct = int(dhcp4[1].strip().split(".")[3])
gw_prefix = util.xml_escape(str(gateway)[:-2])
for ip in range(fist_oct, last_oct + 1):
xml += f"""<host mac='{util.randomMAC()}' ip='{gateway[:-2]}.{ip}' />"""
xml += f"""<host mac='{util.randomMAC()}' ip='{gw_prefix}.{ip}' />"""
xml += """</dhcp>"""
xml += """</ip>"""
if ipv6:
xml += f"""<ip family='ipv6' address='{gateway6}' prefix='{prefix6}'>"""
xml += f"""<ip family='ipv6' address='{esc_gateway6}' prefix='{esc_prefix6}'>"""
if dhcp6:
dhcp6_0 = util.xml_escape(str(dhcp6[0]))
dhcp6_1 = util.xml_escape(str(dhcp6[1]))
xml += f"""<dhcp>
<range start='{dhcp6[0]}' end='{dhcp6[1]}' />"""
<range start='{dhcp6_0}' end='{dhcp6_1}' />"""
xml += """</dhcp>"""
xml += """</ip>"""
xml += """</network>"""
@ -259,10 +272,15 @@ class wvmNetwork(wvmConnect):
parent_index = self.parent_count - 1
for h in hosts:
if h.get("ip") == ip:
h_name = util.xml_escape(h.get("name") or "")
name_attr = f' name="{h_name}"' if h.get("name") else ""
ip_esc = util.xml_escape(ip)
if family == "ipv4":
new_xml = f'<host mac="{h.get("mac")}" name="{h.get("name")}" ip="{ip}"/>'
h_mac = util.xml_escape(h.get("mac") or "")
new_xml = f'<host mac="{h_mac}"{name_attr} ip="{ip_esc}"/>'
if family == "ipv6":
new_xml = f'<host id="{h.get("id")}" name="{h.get("name")}" ip="{ip}"/>'
h_id = util.xml_escape(h.get("id") or "")
new_xml = f'<host id="{h_id}"{name_attr} ip="{ip_esc}"/>'
self.update(
VIR_NETWORK_UPDATE_COMMAND_DELETE,
@ -275,13 +293,17 @@ class wvmNetwork(wvmConnect):
def modify_fixed_address(self, name, address, mac_duid, family="ipv4"):
tree = etree.fromstring(self._XMLDesc(0))
name_esc = util.xml_escape(name) if name else ""
name_attr = f' name="{name_esc}"' if name else ""
mac_esc = util.xml_escape(mac_duid)
ip_addr = ipaddress.ip_address(address)
if family == "ipv4":
new_xml = '<host mac="{}" {} ip="{}"/>'.format(mac_duid, 'name="' + name + '"' if name else "", ipaddress.ip_address(address))
new_xml = f'<host mac="{mac_esc}"{name_attr} ip="{ip_addr}"/>'
hosts = tree.xpath("./ip[not(@family='ipv6')]/dhcp/host")
compare_var = "mac"
parent_index = self.parent_count - 2
if family == "ipv6":
new_xml = '<host id="{}" {} ip="{}"/>'.format(mac_duid, 'name="' + name + '"' if name else "", ipaddress.ip_address(address))
new_xml = f'<host id="{mac_esc}"{name_attr} ip="{ip_addr}"/>'
hosts = tree.xpath("./ip[@family='ipv6']/dhcp/host")
compare_var = "id"
parent_index = self.parent_count - 1
@ -331,8 +353,12 @@ class wvmNetwork(wvmConnect):
return qos_values
def set_qos(self, direction, average, peak, burst):
if direction not in ("inbound","outbound"):
raise Exception("Direction must be inbound or outbound")
if direction not in ("inbound", "outbound"):
raise ValueError("Direction must be inbound or outbound")
average = int(average)
peak = int(peak)
burst = int(burst)
xml = f"<{direction} average='{average}' peak='{peak}' burst='{burst}'/>"
tree = etree.fromstring(self._XMLDesc(0))

92
vrtManager/test_xml_safety.py Normal file
View file

@ -0,0 +1,92 @@
import unittest
from unittest.mock import MagicMock
from django.conf import settings
if not settings.configured:
settings.configure(MAC_OUI="52:54:10")
from vrtManager.instance import wvmInstance
from vrtManager.network import wvmNetworks, wvmNetwork
class TestXmlSafety(unittest.TestCase):
def test_create_network_escaping(self):
wn = wvmNetworks.__new__(wvmNetworks)
wn.wvm = MagicMock()
mock_net = MagicMock()
wn.get_network = MagicMock(return_value=mock_net)
malicious_name = "net<script>alert(1)</script>&name"
malicious_bridge = "br0' onfocus='alert(2)"
wn.create_network(
name=malicious_name,
forward="bridge",
ipv4=False,
gateway=None,
mask=None,
dhcp4=None,
ipv6=False,
gateway6=None,
prefix6=None,
dhcp6=None,
bridge=malicious_bridge,
openvswitch=False,
)
wn.wvm.networkDefineXML.assert_called_once()
defined_xml = wn.wvm.networkDefineXML.call_args[0][0]
self.assertNotIn("<script>", defined_xml)
self.assertIn("&lt;script&gt;alert(1)&lt;/script&gt;&amp;name", defined_xml)
self.assertIn("name='br0&apos; onfocus=&apos;alert(2)'", defined_xml)
def test_network_set_qos_validation(self):
wn = wvmNetwork.__new__(wvmNetwork)
wn._XMLDesc = MagicMock(return_value="<network><name>testnet</name></network>")
wn.wvm = MagicMock()
# Invalid direction
with self.assertRaises(ValueError):
wn.set_qos("lateral", 100, 200, 300)
# Invalid numeric strings
with self.assertRaises(ValueError):
wn.set_qos("inbound", "not_a_number", 200, 300)
def test_snapshot_escaping_and_recovery_on_failure(self):
inst = wvmInstance.__new__(wvmInstance)
inst.get_status = MagicMock(return_value=5) # shutoff
inst.change_snapshot_xml = MagicMock()
inst.recover_snapshot_xml = MagicMock()
inst._XMLDesc = MagicMock(return_value="<domain></domain>")
inst._snapshotCreateXML = MagicMock(side_effect=RuntimeError("libvirt exploded"))
bad_name = "snap<inject>&quot;"
bad_desc = "desc<tag>&"
with self.assertRaises(RuntimeError):
inst.create_snapshot(bad_name, bad_desc)
inst.change_snapshot_xml.assert_called_once()
# Verify recover_snapshot_xml was called in finally despite error
inst.recover_snapshot_xml.assert_called_once()
def test_external_snapshot_recovery(self):
inst = wvmInstance.__new__(wvmInstance)
inst.get_status = MagicMock(return_value=1)
inst.change_snapshot_xml = MagicMock()
inst.recover_snapshot_xml = MagicMock()
inst._XMLDesc = MagicMock(return_value="<domain></domain>")
inst._snapshotCreateXML = MagicMock(return_value=None)
inst.refresh_instance_pools = MagicMock()
inst.create_external_snapshot("ext_snap_1", desc="normal snapshot")
inst.change_snapshot_xml.assert_called_once()
inst.recover_snapshot_xml.assert_called_once()
inst.refresh_instance_pools.assert_called_once()
if __name__ == "__main__":
unittest.main()