use crate::error::{Result, VmSpectError};
use crate::vms::vmdk;
use std::collections::VecDeque;
use std::fs::{self, File};
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
pub const SUPPORTED_EXTENSIONS: &[&str] = &["vmdk", "qcow2", "vdi", "vhd", "vhdx", "raw", "img"];
pub const MAGIC_QCOW2: &[u8; 4] = b"QFI\xfb";
pub const MAGIC_VMDK_KDMV: &[u8; 4] = b"KDMV";
pub const MAGIC_VHDX: &[u8; 8] = b"vhdxfile";
pub const MAGIC_VHD_CONECTIX: &[u8; 8] = b"conectix";
pub const MAGIC_VDI_SIGNATURE: &[u8; 4] = &[0x7F, 0x10, 0xDA, 0xBE];
pub const MAGIC_VDI_PREFIX_SUN: &[u8] = b"<<< Sun VirtualBox Disk Image >>>";
pub const MAGIC_VDI_PREFIX_ORACLE: &[u8] = b"<<< Oracle VM VirtualBox Disk Image >>>";
pub fn is_secondary_extent(path: &Path) -> bool {
let name = match path.file_name().and_then(|n| n.to_str()) {
Some(n) => n.to_ascii_lowercase(),
None => return false,
};
if name.ends_with(".vmdk") {
let stem = match path.file_stem().and_then(|s| s.to_str()) {
Some(s) => s.to_ascii_lowercase(),
None => return false,
};
if stem.ends_with("-flat")
|| stem.ends_with("_flat")
|| stem.ends_with("-delta")
|| stem.ends_with("_delta")
|| stem.ends_with("-sesparse")
|| stem.ends_with("_sesparse")
{
return true;
}
for sep in &["-s", "_s"] {
if let Some(pos) = stem.rfind(sep) {
let suffix = &stem[pos + sep.len()..];
if !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()) {
return true;
}
}
}
} else if name.ends_with(".vhd") || name.ends_with(".vhdx") {
let stem = match path.file_stem().and_then(|s| s.to_str()) {
Some(s) => s.to_ascii_lowercase(),
None => return false,
};
if stem.ends_with("-sys")
|| stem.ends_with("_sys")
|| stem.ends_with("-delta")
|| stem.ends_with("_delta")
{
return true;
}
}
false
}
pub fn is_vm_image(path: &Path) -> bool {
if path.is_dir() {
return false;
}
let ext = match path.extension().and_then(|e| e.to_str()) {
Some(e) => e.to_ascii_lowercase(),
None => return false,
};
if !SUPPORTED_EXTENSIONS.contains(&ext.as_str()) {
return false;
}
if is_secondary_extent(path) {
return false;
}
if let Ok(meta) = fs::metadata(path) {
if !meta.is_file() || meta.len() == 0 {
return false;
}
}
true
}
pub fn list_vms(directory: &Path, recursive: bool) -> Result<Vec<PathBuf>> {
if !directory.exists() {
return Err(VmSpectError::ImageNotFound(format!(
"Directory not found: {}",
directory.display()
)));
}
if !directory.is_dir() {
return Err(VmSpectError::Other(format!(
"The supplied path is not a directory: {}",
directory.display()
)));
}
let mut images = Vec::new();
let mut queue = VecDeque::new();
queue.push_back(directory.to_path_buf());
while let Some(current_dir) = queue.pop_front() {
let entries = match fs::read_dir(¤t_dir) {
Ok(entries) => entries,
Err(e) => return Err(VmSpectError::Io(e)),
};
for entry in entries {
let entry = match entry {
Ok(e) => e,
Err(e) => return Err(VmSpectError::Io(e)),
};
let path = entry.path();
let file_type = match entry.file_type() {
Ok(ft) => ft,
Err(e) => return Err(VmSpectError::Io(e)),
};
if file_type.is_dir() {
if recursive {
queue.push_back(path);
}
} else if file_type.is_file() && is_vm_image(&path) {
images.push(path);
}
}
}
images.sort();
Ok(images)
}
pub fn count_vms(directory: &Path, recursive: bool) -> Result<usize> {
list_vms(directory, recursive).map(|list| list.len())
}
pub fn has_vms(directory: &Path, recursive: bool) -> Result<bool> {
if !directory.exists() {
return Err(VmSpectError::ImageNotFound(format!(
"Directory not found: {}",
directory.display()
)));
}
if !directory.is_dir() {
return Err(VmSpectError::Other(format!(
"The supplied path is not a directory: {}",
directory.display()
)));
}
let mut queue = VecDeque::new();
queue.push_back(directory.to_path_buf());
while let Some(current_dir) = queue.pop_front() {
let entries = match fs::read_dir(¤t_dir) {
Ok(entries) => entries,
Err(e) => return Err(VmSpectError::Io(e)),
};
for entry in entries {
let entry = match entry {
Ok(e) => e,
Err(e) => return Err(VmSpectError::Io(e)),
};
let path = entry.path();
let file_type = match entry.file_type() {
Ok(ft) => ft,
Err(e) => return Err(VmSpectError::Io(e)),
};
if file_type.is_dir() {
if recursive {
queue.push_back(path);
}
} else if file_type.is_file() && is_vm_image(&path) {
return Ok(true);
}
}
}
Ok(false)
}
pub fn verify_image_integrity(path: &Path) -> Result<bool> {
if !path.exists() {
return Err(VmSpectError::ImageNotFound(format!(
"Image not found: {}",
path.display()
)));
}
let mut file = File::open(path).map_err(VmSpectError::Io)?;
let size = file.metadata().map_err(VmSpectError::Io)?.len();
if size == 0 {
return Ok(false);
}
let to_read = (size as usize).min(4096);
let mut header = vec![0u8; to_read];
file.read_exact(&mut header).map_err(VmSpectError::Io)?;
let ext = path
.extension()
.and_then(|e| e.to_str())
.map(|e| e.to_ascii_lowercase())
.unwrap_or_default();
match ext.as_str() {
"qcow2" => Ok(header.len() >= 4 && header.starts_with(MAGIC_QCOW2)),
"vmdk" => {
if header.len() >= 4 && header.starts_with(MAGIC_VMDK_KDMV) {
return Ok(true);
}
let text = String::from_utf8_lossy(&header);
let trimmed_text = text.trim_start();
if trimmed_text.starts_with("# Disk DescriptorFile")
|| trimmed_text.starts_with("# VMDK Header")
|| trimmed_text.starts_with("# VMDK")
|| trimmed_text.contains("# Disk DescriptorFile")
{
return Ok(true);
}
Ok(false)
}
"vdi" => {
if header.starts_with(b"<<< ") {
if header.starts_with(MAGIC_VDI_PREFIX_SUN)
|| header.starts_with(MAGIC_VDI_PREFIX_ORACLE)
{
return Ok(true);
}
if header.len() >= 0x44 && header[0x40..0x44] == *MAGIC_VDI_SIGNATURE {
return Ok(true);
}
}
if header.len() >= 0x44 && header[0x40..0x44] == *MAGIC_VDI_SIGNATURE {
return Ok(true);
}
Ok(false)
}
"vhdx" => Ok(header.len() >= 8 && header.starts_with(MAGIC_VHDX)),
"vhd" => {
if header.len() >= 8
&& (header.starts_with(MAGIC_VHD_CONECTIX) || header.starts_with(b"cxsparse"))
{
return Ok(true);
}
if size >= 512 {
let mut footer = [0u8; 512];
file.seek(SeekFrom::Start(size - 512))
.map_err(VmSpectError::Io)?;
file.read_exact(&mut footer).map_err(VmSpectError::Io)?;
if footer.starts_with(MAGIC_VHD_CONECTIX) {
return Ok(true);
}
}
Ok(false)
}
"raw" | "img" => {
if size < 512 {
return Ok(false);
}
if header.len() >= 512 && header[510] == 0x55 && header[511] == 0xAA {
return Ok(true);
}
if header.len() >= 520 && &header[512..520] == b"EFI PART" {
return Ok(true);
}
Ok(true)
}
_ => {
if header.starts_with(MAGIC_QCOW2)
|| header.starts_with(MAGIC_VMDK_KDMV)
|| header.starts_with(MAGIC_VHDX)
|| header.starts_with(MAGIC_VHD_CONECTIX)
|| (header.len() >= 0x44 && header[0x40..0x44] == *MAGIC_VDI_SIGNATURE)
{
Ok(true)
} else {
Ok(false)
}
}
}
}
pub fn requires_nbd(path: &Path) -> Result<bool> {
if !path.exists() {
return Err(VmSpectError::ImageNotFound(format!(
"Image not found: {}",
path.display()
)));
}
let ext = path
.extension()
.and_then(|e| e.to_str())
.map(|e| e.to_ascii_lowercase())
.unwrap_or_default();
match ext.as_str() {
"raw" | "img" => Ok(false),
"qcow2" | "vdi" | "vhd" | "vhdx" => Ok(true),
"vmdk" => {
let mut file = File::open(path).map_err(VmSpectError::Io)?;
let mut header = [0u8; 512];
let n = file.read(&mut header).map_err(VmSpectError::Io)?;
let header = &header[..n];
if vmdk::is_sparse_header(header) {
let cab = match vmdk::read_sparse_header(header) {
Ok(c) => c,
Err(_) => return Ok(true),
};
if cab.unsupported_reason().is_some() {
return Ok(true);
}
if cab.descriptor_offset != 0 && cab.descriptor_sectors != 0 {
let bytes_desc = (cab.descriptor_sectors * vmdk::SECTOR) as usize;
let mut text = vec![0u8; bytes_desc.min(64 * 1024)];
if file
.seek(SeekFrom::Start(cab.descriptor_offset * vmdk::SECTOR))
.is_ok()
&& file.read_exact(&mut text).is_ok()
{
let d = vmdk::parse_descriptor(&String::from_utf8_lossy(&text));
if d.has_parent() || d.extents.len() > 1 {
return Ok(true);
}
}
}
Ok(false)
} else if vmdk::is_text_descriptor(header) {
let text = fs::read_to_string(path).map_err(VmSpectError::Io)?;
let d = vmdk::parse_descriptor(&text);
if d.has_parent() || d.extents.is_empty() {
return Ok(true);
}
if d.extents.len() > 1
|| d.create_type
.to_ascii_lowercase()
.contains("twogbmaxextent")
{
return Ok(true);
}
if d.extents.len() == 1 {
let kind = d.extents[0].kind.to_ascii_uppercase();
if kind == "FLAT" || kind == "ZERO" {
return Ok(false);
}
}
Ok(true)
} else {
Ok(true)
}
}
_ => Ok(true),
}
}
pub fn requires_qemu(path: &Path) -> Result<bool> {
requires_nbd(path)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::tempdir;
#[test]
fn test_is_secondary_extent() {
assert!(is_secondary_extent(Path::new("ubuntu-flat.vmdk")));
assert!(is_secondary_extent(Path::new("ubuntu_flat.vmdk")));
assert!(is_secondary_extent(Path::new("ubuntu-delta.vmdk")));
assert!(is_secondary_extent(Path::new("ubuntu-sesparse.vmdk")));
assert!(is_secondary_extent(Path::new("windows-s001.vmdk")));
assert!(is_secondary_extent(Path::new("windows-s02.vmdk")));
assert!(is_secondary_extent(Path::new("windows_s1.vmdk")));
assert!(is_secondary_extent(Path::new("disk-sys.vhd")));
assert!(is_secondary_extent(Path::new("disk_sys.vhd")));
assert!(is_secondary_extent(Path::new("disk-delta.vhd")));
assert!(is_secondary_extent(Path::new("disk-delta.vhdx")));
assert!(!is_secondary_extent(Path::new("ubuntu.vmdk")));
assert!(!is_secondary_extent(Path::new("windows-server.vmdk")));
assert!(!is_secondary_extent(Path::new("disk.qcow2")));
assert!(!is_secondary_extent(Path::new("disk.vdi")));
assert!(!is_secondary_extent(Path::new("disk.vhd")));
assert!(!is_secondary_extent(Path::new("disk.vhdx")));
assert!(!is_secondary_extent(Path::new("disk.raw")));
}
#[test]
fn test_is_vm_image() {
assert!(is_vm_image(Path::new("vm.vmdk")));
assert!(is_vm_image(Path::new("vm.qcow2")));
assert!(is_vm_image(Path::new("vm.vdi")));
assert!(is_vm_image(Path::new("vm.vhd")));
assert!(is_vm_image(Path::new("vm.vhdx")));
assert!(is_vm_image(Path::new("vm.raw")));
assert!(is_vm_image(Path::new("vm.img")));
assert!(!is_vm_image(Path::new("vm-flat.vmdk")));
assert!(!is_vm_image(Path::new("vm-s001.vmdk")));
assert!(!is_vm_image(Path::new("vm.iso")));
assert!(!is_vm_image(Path::new("vm.txt")));
}
#[test]
fn test_verify_integrity_qcow2() {
let dir = tempdir().unwrap();
let path = dir.path().join("test.qcow2");
let mut f = File::create(&path).unwrap();
f.write_all(b"QFI\xfb\x00\x00\x00\x03").unwrap();
assert!(verify_image_integrity(&path).unwrap());
let invalid_path = dir.path().join("invalid.qcow2");
let mut f2 = File::create(&invalid_path).unwrap();
f2.write_all(b"NOT_QCOW2_HEADER").unwrap();
assert!(!verify_image_integrity(&invalid_path).unwrap());
}
#[test]
fn test_verify_integrity_vmdk() {
let dir = tempdir().unwrap();
let sparse_path = dir.path().join("sparse.vmdk");
let mut f1 = File::create(&sparse_path).unwrap();
f1.write_all(b"KDMV\x01\x00\x00\x00").unwrap();
assert!(verify_image_integrity(&sparse_path).unwrap());
let desc_path = dir.path().join("desc.vmdk");
let mut f2 = File::create(&desc_path).unwrap();
f2.write_all(b"# Disk DescriptorFile\nversion=1\nCID=fffffffe\n")
.unwrap();
assert!(verify_image_integrity(&desc_path).unwrap());
}
#[test]
fn test_verify_integrity_vdi_vhdx_vhd() {
let dir = tempdir().unwrap();
let vhdx_path = dir.path().join("test.vhdx");
let mut f = File::create(&vhdx_path).unwrap();
f.write_all(b"vhdxfile\x00\x00\x00\x00").unwrap();
assert!(verify_image_integrity(&vhdx_path).unwrap());
let vhd_path = dir.path().join("test.vhd");
let mut f2 = File::create(&vhd_path).unwrap();
f2.write_all(b"conectix\x00\x00\x00\x00").unwrap();
assert!(verify_image_integrity(&vhd_path).unwrap());
let vdi_path = dir.path().join("test.vdi");
let mut f3 = File::create(&vdi_path).unwrap();
let mut header_vdi = vec![0u8; 100];
header_vdi[..MAGIC_VDI_PREFIX_ORACLE.len()].copy_from_slice(MAGIC_VDI_PREFIX_ORACLE);
header_vdi[0x40..0x44].copy_from_slice(MAGIC_VDI_SIGNATURE);
f3.write_all(&header_vdi).unwrap();
assert!(verify_image_integrity(&vdi_path).unwrap());
}
#[test]
fn test_list_count_has_vms() {
let dir = tempdir().unwrap();
let sub = dir.path().join("subdir");
fs::create_dir(&sub).unwrap();
let vm1 = dir.path().join("ubuntu.qcow2");
let vm2 = dir.path().join("disk.vmdk");
let extent = dir.path().join("disk-flat.vmdk");
let vm3 = sub.join("windows.vhdx");
let dummy = dir.path().join("notes.txt");
File::create(&vm1).unwrap().write_all(b"data").unwrap();
File::create(&vm2).unwrap().write_all(b"data").unwrap();
File::create(&extent).unwrap().write_all(b"data").unwrap();
File::create(&vm3).unwrap().write_all(b"data").unwrap();
File::create(&dummy).unwrap().write_all(b"data").unwrap();
let vms_non_rec = list_vms(dir.path(), false).unwrap();
assert_eq!(vms_non_rec.len(), 2);
assert!(vms_non_rec.contains(&vm1));
assert!(vms_non_rec.contains(&vm2));
assert!(!vms_non_rec.contains(&extent));
assert_eq!(count_vms(dir.path(), false).unwrap(), 2);
assert!(has_vms(dir.path(), false).unwrap());
let vms_rec = list_vms(dir.path(), true).unwrap();
assert_eq!(vms_rec.len(), 3);
assert!(vms_rec.contains(&vm3));
assert_eq!(count_vms(dir.path(), true).unwrap(), 3);
assert!(has_vms(dir.path(), true).unwrap());
let empty_dir = tempdir().unwrap();
assert_eq!(count_vms(empty_dir.path(), true).unwrap(), 0);
assert!(!has_vms(empty_dir.path(), true).unwrap());
}
#[test]
fn test_requires_nbd_formats() {
let dir = tempdir().unwrap();
let raw = dir.path().join("disk.raw");
File::create(&raw).unwrap().write_all(b"raw data").unwrap();
assert!(!requires_nbd(&raw).unwrap());
let qcow2 = dir.path().join("disk.qcow2");
File::create(&qcow2).unwrap().write_all(b"qcow2").unwrap();
assert!(requires_nbd(&qcow2).unwrap());
assert!(requires_qemu(&qcow2).unwrap());
let vhdx = dir.path().join("disk.vhdx");
File::create(&vhdx).unwrap().write_all(b"vhdx").unwrap();
assert!(requires_nbd(&vhdx).unwrap());
let vmdk_flat_desc = dir.path().join("monolithic_flat.vmdk");
let mut f_desc = File::create(&vmdk_flat_desc).unwrap();
f_desc
.write_all(
b"# Disk DescriptorFile\ncreateType=\"monolithicFlat\"\nRW 2048 FLAT \"data.flat\" 0\n",
)
.unwrap();
assert!(!requires_nbd(&vmdk_flat_desc).unwrap());
let vmdk_snap = dir.path().join("snapshot.vmdk");
let mut f_snap = File::create(&vmdk_snap).unwrap();
f_snap
.write_all(
b"# Disk DescriptorFile\nparentFileNameHint=\"base.vmdk\"\nRW 2048 FLAT \"snap.flat\" 0\n",
)
.unwrap();
assert!(requires_nbd(&vmdk_snap).unwrap());
}
}