ambient-ci 0.14.0

//! Manage a collection of virtual machine images.

use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};

use serde::{Deserialize, Serialize};

use crate::{
    checksummer::Checksummer,
    util::{cat_text_file, mkdir, now, write_file, UtilError},
};

/// Build metadata for an image.
pub struct MetadataBuilder {
    source: PathBuf,
    name: String,
    filename: PathBuf,
    sha256: String,
    description: Option<String>,
    url: Option<String>,
    import_date: String,
    uefi: bool,
}

impl MetadataBuilder {
    /// Create a new metadata builder.
    pub fn new(name: &str, source: &Path) -> Result<Self, ImageStoreError> {
        let imported_image = PathBuf::from(format!("{name}.qcow2"));

        let sha256 = Checksummer::new(source)
            .sha256()
            .map_err(|err| ImageStoreError::Digest(source.into(), err))?
            .to_string();

        Ok(Self {
            source: source.into(),
            name: name.into(),
            filename: imported_image,
            sha256,
            description: None,
            url: None,
            import_date: now()?,
            uefi: false,
        })
    }

    /// Set description of image.
    pub fn description(&mut self, value: &str) {
        self.description = Some(value.into());
    }

    /// Set URL from whence image came.
    pub fn url(&mut self, value: &str) {
        self.url = Some(value.into());
    }

    /// Does image require UEFI?
    pub fn uefi(&mut self, value: bool) {
        self.uefi = value;
    }

    /// Build metadata.
    pub fn build(self) -> Metadata {
        Metadata {
            source: self.source,
            name: self.name,
            filename: self.filename,
            sha256: self.sha256,
            description: self.description,
            url: self.url,
            import_date: self.import_date,
            uefi: self.uefi,
        }
    }
}

/// Metadata for a virtual machine image.
#[derive(Debug, Serialize, Deserialize)]
pub struct Metadata {
    // File from which image was copied into image store.
    #[serde(skip)]
    source: PathBuf,

    /// User-supplied Name of image.
    pub name: String,

    /// Filename relative to image store.
    pub filename: PathBuf,

    /// SHA256 checksum of image file.
    pub sha256: String,

    /// A user-supplied description of image.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// URL to origin of image.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub url: Option<String>,

    /// Date when image was imported.
    pub import_date: String,

    /// Boot image with UEFI?
    pub uefi: bool,
}

impl Metadata {
    /// Serialize metadata to JSON.
    pub fn to_json(&self) -> Result<String, ImageStoreError> {
        let json: String = serde_json::to_string_pretty(&self)
            .map_err(|err| ImageStoreError::JsonImageSer(self.name.clone(), err))?;
        Ok(json)
    }

    /// Does image require UEFI?
    pub fn uefi(&self) -> bool {
        self.uefi
    }
}

/// A collection of virtual machine images.
#[derive(Debug)]
pub struct ImageStore {
    dirname: PathBuf,
    metadata: HashMap<String, Metadata>,
}

impl ImageStore {
    /// Create a new [`ImageStore`] that refers to a directory.
    pub fn new(dirname: &Path) -> Result<Self, ImageStoreError> {
        let filename = Self::metadata_filename(dirname);
        let metadata = if filename.exists() {
            let json = cat_text_file(&filename).map_err(ImageStoreError::Util)?;
            serde_json::from_str(&json).map_err(|err| ImageStoreError::JsonParse(filename, err))?
        } else {
            HashMap::new()
        };

        Ok(Self {
            dirname: dirname.into(),
            metadata,
        })
    }

    /// Save changes to disk.
    pub fn save(&self) -> Result<(), ImageStoreError> {
        let filename = Self::metadata_filename(&self.dirname);
        let json: String = serde_json::to_string(&self.metadata)
            .map_err(|err| ImageStoreError::JsonSer(filename.clone(), err))?;
        write_file(&filename, json.as_bytes()).map_err(ImageStoreError::Util)?;
        Ok(())
    }

    /// Does the store contain an image with a given name?
    pub fn contains(&self, name: &str) -> bool {
        self.metadata.contains_key(name)
    }

    fn metadata_filename(dirname: &Path) -> PathBuf {
        dirname.join("images.json")
    }

    /// Name of image file.
    pub fn image_filename(&self, metadata: &Metadata) -> PathBuf {
        self.dirname.join(&metadata.filename)
    }

    /// List names of all images.
    pub fn image_names(&self) -> Result<Vec<String>, ImageStoreError> {
        Ok(self.metadata.keys().map(|name| name.to_string()).collect())
    }

    /// Metadata for an image.
    pub fn get_metadata(&self, name: &str) -> Option<&Metadata> {
        self.metadata.get(name)
    }

    /// Import an image to the store.
    pub fn import(&mut self, metadata: Metadata) -> Result<(), ImageStoreError> {
        if !self.dirname.exists() {
            mkdir(&self.dirname).map_err(ImageStoreError::Util)?;
        }

        let filename = self.dirname.join(&metadata.filename);
        std::fs::copy(&metadata.source, &filename)
            .map_err(|err| ImageStoreError::Copy(filename.clone(), err))?;

        self.metadata.insert(metadata.name.clone(), metadata);

        Ok(())
    }

    /// Remove an image from a store.
    pub fn remove(&mut self, name: &str) -> Result<(), ImageStoreError> {
        let filename = self.image_name(name);
        std::fs::remove_file(&filename)
            .map_err(|err| ImageStoreError::RemoveImageFile(filename, err))?;
        self.metadata.remove(name);
        Ok(())
    }

    fn image_name(&self, name: &str) -> PathBuf {
        self.dirname.join(format!("{name}.qcow2"))
    }
}

/// Possible errors from managing a set of virtual machine images.
#[derive(Debug, thiserror::Error)]
pub enum ImageStoreError {
    /// Can't compute checksum.
    #[error("failed to compute checksum for file {0}")]
    Digest(PathBuf, #[source] crate::checksummer::ChecksummerError),

    /// Can't image store metadata as JSON.
    #[error("failed to parse image store metadata file {0} as JSON")]
    JsonParse(PathBuf, #[source] serde_json::Error),

    /// Can't encode image store metadata as JSON.
    #[error("failed to encode image store {0} metadata as JSON")]
    JsonSer(PathBuf, #[source] serde_json::Error),

    /// Can't encodce image metadata as JSON.
    #[error("failed to encode image {0} metadata as JSON")]
    JsonImageSer(String, #[source] serde_json::Error),

    /// Error from `util` module.
    #[error(transparent)]
    Util(#[from] UtilError),

    /// Can't copy image file.
    #[error("failed to copy image {0} into image store")]
    Copy(PathBuf, #[source] std::io::Error),

    /// Can't remove image file.
    #[error("failed to remove image file {0}")]
    RemoveImageFile(PathBuf, #[source] std::io::Error),
}