Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
146 changes: 79 additions & 67 deletions lib/cuckoo/common/abstracts.py
Original file line number Diff line number Diff line change
Expand Up @@ -437,6 +437,8 @@ def __init__(self):
)

super().__init__()
self.conn = None
self.conn_lock = threading.Lock()

def _initialize_check(self):
"""Runs all checks when a machine manager is initialized.
Expand All @@ -446,9 +448,6 @@ def _initialize_check(self):
if not self._version_check():
raise CuckooMachineError("Libvirt version is not supported, please get an updated version")

# Preload VMs
self.vms = self._fetch_machines()

# Base checks. Also attempts to shutdown any machines which are
# currently still active.
super()._initialize_check()
Expand All @@ -474,34 +473,33 @@ def start(self, label=None):

conn = self._connect(label)

snapshot_list = self.vms[label].snapshotListNames(flags=0)
try:
vm = conn.lookupByName(label)
except libvirt.libvirtError as e:
raise CuckooMachineError(f"Cannot find machine {label}") from e

# If a snapshot is configured try to use it.
if vm_info.snapshot and vm_info.snapshot in snapshot_list:
# Revert to desired snapshot, if it exists.
log.debug("Using snapshot %s for virtual machine %s", vm_info.snapshot, label)
try:
vm = self.vms[label]
snapshot = None
try:
snapshot_list = vm.snapshotListNames(flags=0)

# If a snapshot is configured try to use it.
if vm_info.snapshot and vm_info.snapshot in snapshot_list:
log.debug("Using snapshot %s for virtual machine %s", vm_info.snapshot, label)
snapshot = vm.snapshotLookupByName(vm_info.snapshot, flags=0)
self.vms[label].revertToSnapshot(snapshot, flags=0)
except libvirt.libvirtError as e:
msg = f"Unable to restore snapshot {vm_info.snapshot} on virtual machine {label}. Your snapshot MUST BE in running state!"
raise CuckooMachineError(msg) from e
finally:
self._disconnect(conn)
elif self._get_snapshot(label):
snapshot = self._get_snapshot(label)
log.debug("Using snapshot %s for virtual machine %s", snapshot.getName(), label)
try:
self.vms[label].revertToSnapshot(snapshot, flags=0)
except libvirt.libvirtError as e:
raise CuckooMachineError(f"Unable to restore snapshot on virtual machine {label}. Your snapshot MUST BE in running state!") from e
finally:
self._disconnect(conn)
else:
self._disconnect(conn)
else:
snapshot = self._get_snapshot(label, vm)
except libvirt.libvirtError:
log.warning("Unable to fetch snapshot for virtual machine %s", label)

if not snapshot:
raise CuckooMachineError(f"No snapshot found for virtual machine {label}")

log.debug("Using snapshot %s for virtual machine %s", snapshot.getName(), label)
try:
vm.revertToSnapshot(snapshot, flags=0)
except libvirt.libvirtError as e:
raise CuckooMachineError(f"Unable to restore snapshot on virtual machine {label}. Your snapshot MUST BE in running state!") from e

# Check state.
self._wait_status(label, self.RUNNING)

Expand All @@ -521,14 +519,14 @@ def stop(self, label=None):
# Force virtual machine shutdown.
conn = self._connect(label)
try:
if not self.vms[label].isActive():
vm = conn.lookupByName(label)
if not vm.isActive():
log.debug("Trying to stop an already stopped machine %s, skipping", label)
else:
self.vms[label].destroy() # Machete's way!
vm.destroy() # Machete's way!
except libvirt.libvirtError as e:
raise CuckooMachineError(f"Error stopping virtual machine {label}: {e}") from e
finally:
self._disconnect(conn)

# Check state.
self._wait_status(label, self.POWEROFF)

Expand All @@ -544,8 +542,14 @@ def shutdown(self):
except CuckooMachineError as e:
log.warning("Unable to shutdown machine %s, please check manually. Error: %s", machine.label, e)

# Free handlers.
self.vms = None
# Close connection
with self.conn_lock:
if self.conn:
try:
self.conn.close()
except libvirt.libvirtError:
pass
self.conn = None

def screenshot(self, label, path):
"""Screenshot a running virtual machine.
Expand Down Expand Up @@ -587,11 +591,10 @@ def dump_memory(self, label, path):
# it'll still be owned by root, so we can't delete it, but at least we can read it
fd = open(path, "w")
fd.close()
self.vms[label].coreDump(path, flags=libvirt.VIR_DUMP_MEMORY_ONLY)
vm = conn.lookupByName(label)
vm.coreDump(path, flags=libvirt.VIR_DUMP_MEMORY_ONLY)
except libvirt.libvirtError as e:
raise CuckooMachineError(f"Error dumping memory virtual machine {label}: {e}") from e
finally:
self._disconnect(conn)

def _status(self, label):
"""Gets current status of a vm.
Expand All @@ -613,11 +616,10 @@ def _status(self, label):

conn = self._connect(label)
try:
state = self.vms[label].state(flags=0)
vm = conn.lookupByName(label)
state = vm.state(flags=0)
except libvirt.libvirtError as e:
raise CuckooMachineError(f"Error getting status for virtual machine {label}: {e}") from e
finally:
self._disconnect(conn)

if state:
if state[0] == 1:
Expand All @@ -644,29 +646,37 @@ def _connect(self, label=None):
if not self.dsn:
raise CuckooMachineError("You must provide a proper connection string")

try:
return libvirt.open(self.dsn)
except libvirt.libvirtError as e:
raise CuckooMachineError("Cannot connect to libvirt") from e
with self.conn_lock:
if self.conn:
try:
if self.conn.isAlive():
return self.conn
except libvirt.libvirtError:
pass

# Connection is dead
try:
self.conn.close()
except libvirt.libvirtError:
pass
self.conn = None

try:
self.conn = libvirt.open(self.dsn)
except libvirt.libvirtError as e:
raise CuckooMachineError("Cannot connect to libvirt") from e

def _disconnect(self, conn):
return self.conn

def _disconnect(self, _conn):
"""Disconnects to libvirt subsystem.
@raise CuckooMachineError: if cannot disconnect from libvirt.
"""
try:
conn.close()
except libvirt.libvirtError as e:
raise CuckooMachineError("Cannot disconnect from libvirt") from e

def _fetch_machines(self):
"""Fetch machines handlers.
@return: dict with machine label as key and handle as value.
"""
return {vm.label: self._lookup(vm.label) for vm in self.machines()}
# Do nothing, keep connection open for reuse
pass

def _lookup(self, label):
"""Search for a virtual machine.
@param conn: libvirt connection handle.
@param label: virtual machine name.
@raise CuckooMachineError: if virtual machine is not found.
"""
Expand All @@ -675,8 +685,6 @@ def _lookup(self, label):
vm = conn.lookupByName(label)
except libvirt.libvirtError as e:
raise CuckooMachineError(f"Cannot find machine {label}") from e
finally:
self._disconnect(conn)
return vm

def _list(self):
Expand All @@ -685,22 +693,31 @@ def _list(self):
"""
conn = self._connect()
try:
if hasattr(conn, "listAllDomains"):
# flags=0 returns all domains (active and inactive)
return [dom.name() for dom in conn.listAllDomains(0)]

# Fallback for older libvirt versions
names = conn.listDefinedDomains()
for vid in conn.listDomainsID():
try:
dom = conn.lookupByID(vid)
names.append(dom.name())
except libvirt.libvirtError:
continue
return names
except libvirt.libvirtError as e:
raise CuckooMachineError("Cannot list domains") from e
finally:
self._disconnect(conn)
return names

def _version_check(self):
"""Check if libvirt release supports snapshots.
@return: True or false.
"""
return libvirt.getVersion() >= 8000

def _get_snapshot(self, label):
def _get_snapshot(self, label, vm):
"""Get current snapshot for virtual machine
@param label: virtual machine name
@param vm: virtual machine handle
@return None or current snapshot
@raise CuckooMachineError: if cannot find current snapshot or
when there are too many snapshots available
Expand All @@ -715,10 +732,7 @@ def _extract_creation_time(node):
return xml.findtext("./creationTime")

snapshot = None
conn = self._connect(label)
try:
vm = self.vms[label]

# Try to get the currrent snapshot, otherwise fallback on the latest
# from config file.
if vm.hasCurrentSnapshot(flags=0):
Expand All @@ -732,8 +746,6 @@ def _extract_creation_time(node):
snapshot = sorted(all_snapshots, key=_extract_creation_time, reverse=True)[0]
except libvirt.libvirtError:
raise CuckooMachineError(f"Unable to get snapshot for virtual machine {label}")
finally:
self._disconnect(conn)

return snapshot

Expand Down
6 changes: 3 additions & 3 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2119,9 +2119,9 @@ pyopenssl==25.0.0 ; python_version >= "3.10" and python_version < "4.0" \
pyparsing==3.2.1 ; python_version >= "3.10" and python_version < "4.0" \
--hash=sha256:506ff4f4386c4cec0590ec19e6302d3aedb992fdc02c761e90416f158dacf8e1 \
--hash=sha256:61980854fd66de3a90028d679a954d5f2623e83144b5afe5ee86f43d762e5f0a
pypdf==5.2.0 ; python_version >= "3.10" and python_version < "4.0" \
--hash=sha256:7c38e68420f038f2c4998fd9d6717b6db4f6cef1642e9cf384d519c9cf094663 \
--hash=sha256:d107962ec45e65e3bd10c1d9242bdbbedaa38193c9e3a6617bd6d996e5747b19
pypdf==6.6.2 ; python_version >= "3.10" and python_version < "4.0" \
--hash=sha256:0a3ea3b3303982333404e22d8f75d7b3144f9cf4b2970b96856391a516f9f016 \
--hash=sha256:44c0c9811cfb3b83b28f1c3d054531d5b8b81abaedee0d8cb403650d023832ba
pyre2==0.3.10 ; python_version >= "3.10" and python_version < "4.0" \
--hash=sha256:2771bc0a3a5f3fd1d34fe8ae80debd1fe59a1cdcecdbeb60818bff8deb247e56 \
--hash=sha256:29312fbb22b7d3cf3e522c43267098162e98ec8b4fc202c9260e22df671b42e1 \
Expand Down