ambient-ci 0.14.0

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

use clap::Parser;
use tempfile::tempdir;

use ambient_ci::{
    action::UnsafeAction,
    cloud_init::{CloudInitError, LocalDataStoreBuilder},
    image_store::{ImageStore, ImageStoreError, MetadataBuilder},
    plan::RunnablePlan,
    qemu::{self, QemuError, QemuRunner},
    qemu_utils::{convert_image, QemuUtilError},
    run::{create_cloud_init_iso, create_executor_vdrive, create_source_vdrive},
    runlog::{RunLog, RunLogSource},
    util::{cat_text_file, mkdir, UtilError},
    vdrive::VirtualDriveError,
};

use super::{AmbientError, Config, Leaf};

const SECRET: &str = "xyzzy";

/// Prepare image.
#[derive(Debug, Parser)]
pub struct PrepareImage {
    /// The name of the image in the image store to use as the base image.
    #[clap(long)]
    base: String,

    /// The name of the image in the new prepared image in the image store.
    #[clap(long)]
    new: String,

    /// Location of the ambient-execute-plan binary.
    #[clap(long)]
    executor: PathBuf,

    /// Enable network?
    #[clap(long)]
    network: bool,

    /// Run log goes to this file.
    #[clap(long)]
    run_log: Option<PathBuf>,

    /// Console log goes to this file.
    #[clap(long)]
    console_log: Option<PathBuf>,

    /// Run shell command to modify image.
    #[clap(long)]
    shell: Option<String>,
}

impl Leaf for PrepareImage {
    fn run(&self, config: &Config, _runlog: &mut RunLog) -> Result<(), AmbientError> {
        let mut image_store =
            ImageStore::new(config.image_store()).map_err(ImageError::ImageStore)?;
        if image_store.contains(&self.new) {
            return Err(ImageError::ExistsAlready(self.new.clone()).into());
        }
        if let Some(metadata) = image_store.get_metadata(&self.base) {
            let filename = image_store.image_filename(metadata);
            if !filename.exists() {
                return Err(ImageError::NoSuchFile(filename.clone()).into());
            }

            let filename = std::fs::canonicalize(&filename)
                .map_err(|err| ImageError::Canonicalize(filename.clone(), err))?;

            let tmp = tempdir().map_err(ImageError::TempDir)?;
            let cow_image = tmp.path().join("new.qcow2");
            let src = tmp.path().join("src");
            mkdir(&src).map_err(ImageError::Util)?;

            let mut plan = RunnablePlan::default();

            let shell = self.shell.clone().unwrap_or("echo hello, world".into());

            plan.push_unsafe_actions(
                [
                    UnsafeAction::mkdir(Path::new(qemu::WORKSPACE_DIR)),
                    UnsafeAction::mkdir(Path::new(qemu::SOURCE_DIR)),
                    UnsafeAction::shell(&shell),
                ]
                .iter(),
            );
            plan.set_executor_drive(qemu::EXECUTOR_DRIVE);
            plan.set_source_drive(qemu::SOURCE_DRIVE);
            plan.set_source_dir(qemu::SOURCE_DIR);
            let executor_drive = create_executor_vdrive(&tmp, &plan, &self.executor)
                .map_err(ImageError::CreateDrive)?;
            let source_drive = create_source_vdrive(&tmp, &src).map_err(ImageError::CreateDrive)?;

            let run_log = self
                .run_log
                .clone()
                .unwrap_or_else(|| tmp.path().join("run.log"));
            let console_log = self
                .console_log
                .clone()
                .unwrap_or_else(|| tmp.path().join("console.log"));

            let ds = create_cloud_init_iso(self.network).map_err(ImageError::CreateDrive)?;

            let mut runlog = RunLog::default();
            let res = QemuRunner::default()
                .config(config)
                .base_image(&filename)
                .cow_image(&cow_image)
                .executor(&executor_drive)
                .source(&source_drive)
                .cloud_init(&ds)
                .console_log(&console_log)
                .raw_log(&run_log)
                .network(self.network)
                .run(RunLogSource::Plan, &mut runlog)
                .map_err(ImageError::Qemu);

            res?;

            let full_image = tmp.path().join("full.qcow2");
            convert_image(&cow_image, &full_image).map_err(|err| {
                ImageError::ConvertImage(cow_image.clone(), full_image.clone(), err)
            })?;

            let mut new_metadata =
                MetadataBuilder::new(&self.new, &full_image).map_err(ImageError::ImageStore)?;
            new_metadata.uefi(metadata.uefi());
            let new_metadata = new_metadata.build();

            image_store
                .import(new_metadata)
                .map_err(ImageError::ImageStore)?;
            image_store.save().map_err(ImageError::ImageStore)?;

            Ok(())
        } else {
            Err(ImageError::NoSuchImage(self.base.clone()).into())
        }
    }
}

/// Verify that a virtual machine image is acceptable for use with
/// ambient. This is done by booting a VM using the image, and
/// making sure the VM run the provided executable and then shuts down.
///
/// If the verification doesn't finish within a reasonable time, the
/// image is not acceptable.
#[derive(Debug, Parser)]
pub struct VerifyImage {
    /// The name of the image in the image store to test.
    #[clap(long)]
    name: String,

    /// Location of the ambient-execute-plan binary.
    #[clap(long)]
    executor: PathBuf,

    /// Enable network?
    #[clap(long)]
    network: bool,

    /// Run log goes to this file.
    #[clap(long)]
    run_log: Option<PathBuf>,

    /// Console log goes to this file.
    #[clap(long)]
    console_log: Option<PathBuf>,
}

impl Leaf for VerifyImage {
    fn run(&self, config: &Config, _runlog: &mut RunLog) -> Result<(), AmbientError> {
        let image_store = ImageStore::new(config.image_store()).map_err(ImageError::ImageStore)?;
        if let Some(metadata) = image_store.get_metadata(&self.name) {
            let filename = image_store.image_filename(metadata);
            if !filename.exists() {
                return Err(ImageError::NoSuchFile(filename.clone()).into());
            }

            let filename = std::fs::canonicalize(&filename)
                .map_err(|err| ImageError::Canonicalize(filename.clone(), err))?;

            let tmp = tempdir().map_err(ImageError::TempDir)?;
            let src = tmp.path().join("src");
            mkdir(&src).map_err(ImageError::Util)?;

            let mut plan = RunnablePlan::default();
            plan.push_unsafe_actions(
                [
                    UnsafeAction::mkdir(Path::new(qemu::WORKSPACE_DIR)),
                    UnsafeAction::mkdir(Path::new(qemu::SOURCE_DIR)),
                    UnsafeAction::shell(&format!("echo {SECRET}")),
                ]
                .iter(),
            );
            plan.set_executor_drive(qemu::EXECUTOR_DRIVE);
            plan.set_source_drive(qemu::SOURCE_DRIVE);
            plan.set_source_dir(qemu::SOURCE_DIR);
            let executor_drive = create_executor_vdrive(&tmp, &plan, &self.executor)
                .map_err(ImageError::CreateDrive)?;
            let source_drive = create_source_vdrive(&tmp, &src).map_err(ImageError::CreateDrive)?;

            let run_log = self
                .run_log
                .clone()
                .unwrap_or_else(|| tmp.path().join("run.log"));
            let console_log = self
                .console_log
                .clone()
                .unwrap_or_else(|| tmp.path().join("console.log"));

            let ds = create_cloud_init_iso(self.network).map_err(ImageError::CreateDrive)?;

            let mut runlog = RunLog::default();
            let res = QemuRunner::default()
                .config(config)
                .base_image(&filename)
                .executor(&executor_drive)
                .source(&source_drive)
                .cloud_init(&ds)
                .console_log(&console_log)
                .raw_log(&run_log)
                .network(self.network)
                .run(RunLogSource::Plan, &mut runlog)
                .map_err(ImageError::Qemu);

            let log = cat_text_file(&run_log).map_err(ImageError::Util)?;
            if !log.contains(SECRET) {
                return Err(ImageError::NotAcceptable(filename.clone()).into());
            }
            println!("image {} is acceptable to Ambient", filename.display());

            res?;
            Ok(())
        } else {
            Err(ImageError::NoSuchImage(self.name.clone()).into())
        }
    }
}

/// Create a cloud-init local data store ISO file.
#[derive(Debug, Parser)]
pub struct CloudInit {
    /// The name of the ISO file.
    #[clap(long)]
    filename: PathBuf,

    /// Commands to run via cloud-init.
    #[clap(long)]
    runcmd: Vec<String>,

    /// Enable networking?
    #[clap(long)]
    network: bool,
}

impl Leaf for CloudInit {
    fn run(&self, _config: &Config, _runlog: &mut RunLog) -> Result<(), AmbientError> {
        let mut ds = LocalDataStoreBuilder::default()
            .with_network(self.network)
            .with_hostname("ambient")
            .with_bootcmd("echo xyzzy2")
            .with_bootcmd("echo more bootcmd");
        for runcmd in self.runcmd.iter() {
            ds = ds.with_runcmd(runcmd);
        }
        let ds = ds.build().map_err(ImageError::CloudInit)?;
        ds.iso(&self.filename).map_err(ImageError::CloudInit)?;
        Ok(())
    }
}

/// List names of images in the image store.
#[derive(Debug, Parser)]
pub struct ListImages {}

impl Leaf for ListImages {
    fn run(&self, config: &Config, _runlog: &mut RunLog) -> Result<(), AmbientError> {
        let image_store = ImageStore::new(config.image_store()).map_err(ImageError::ImageStore)?;
        for name in image_store.image_names().map_err(ImageError::ImageStore)? {
            println!("{name}");
        }
        Ok(())
    }
}

/// Import images in the image store.
#[derive(Debug, Parser)]
pub struct ImportImage {
    /// Name of image in the store.
    name: String,

    /// File name of the image to import and copy into the store.
    image: PathBuf,

    /// Description of image.
    #[clap(short, long)]
    description: Option<String>,

    /// URL for origin of image.
    #[clap(short, long)]
    url: Option<String>,

    /// Should the image be booted with UEFI?
    #[clap(long)]
    uefi: bool,
}

impl Leaf for ImportImage {
    fn run(&self, config: &Config, _runlog: &mut RunLog) -> Result<(), AmbientError> {
        let mut metadata =
            MetadataBuilder::new(&self.name, &self.image).map_err(ImageError::ImageStore)?;
        if let Some(d) = &self.description {
            metadata.description(d);
        }
        if let Some(u) = &self.url {
            metadata.url(u);
        }
        metadata.uefi(self.uefi || config.uefi());
        let metadata = metadata.build();

        let mut image_store =
            ImageStore::new(config.image_store()).map_err(ImageError::ImageStore)?;
        image_store
            .import(metadata)
            .map_err(ImageError::ImageStore)?;
        image_store.save().map_err(ImageError::ImageStore)?;
        Ok(())
    }
}

/// Remove images from the image store.
#[derive(Debug, Parser)]
pub struct RemoveImages {
    /// Names of image to remove.
    images: Vec<String>,
}

impl Leaf for RemoveImages {
    fn run(&self, config: &Config, _runlog: &mut RunLog) -> Result<(), AmbientError> {
        let mut image_store =
            ImageStore::new(config.image_store()).map_err(ImageError::ImageStore)?;
        for image in self.images.iter() {
            image_store.remove(image).map_err(ImageError::ImageStore)?;
        }
        image_store.save().map_err(ImageError::ImageStore)?;
        Ok(())
    }
}

/// Show information about an image in the image store.
#[derive(Debug, Parser)]
pub struct ShowImage {
    /// Name of image.
    image: String,
}

impl Leaf for ShowImage {
    fn run(&self, config: &Config, _runlog: &mut RunLog) -> Result<(), AmbientError> {
        let image_store = ImageStore::new(config.image_store()).map_err(ImageError::ImageStore)?;
        if let Some(image) = image_store.get_metadata(&self.image) {
            println!("{}", image.to_json().map_err(ImageError::ImageStore)?);
            Ok(())
        } else {
            Err(ImageError::NoSuchImage(self.image.clone()))?
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum ImageError {
    #[error(transparent)]
    Qemu(#[from] QemuError),

    #[error("failed to create a temporary directory")]
    TempDir(#[source] std::io::Error),

    #[error(transparent)]
    Util(#[from] UtilError),

    #[error(transparent)]
    VDrive(#[from] VirtualDriveError),

    #[error(transparent)]
    CloudInit(#[from] CloudInitError),

    #[error("image {0} does not exist")]
    NoSuchFile(PathBuf),

    #[error("image {0} does not exist in the image store")]
    NoSuchImage(String),

    #[error("image {0} is NOT an acceptable image for Ambient")]
    NotAcceptable(PathBuf),

    #[error("image store already contains {0}")]
    ExistsAlready(String),

    #[error(transparent)]
    CreateDrive(#[from] ambient_ci::run::RunError),

    #[error(transparent)]
    ImageStore(#[from] ImageStoreError),

    #[error("failed to make filename canonical: {0}")]
    Canonicalize(PathBuf, #[source] std::io::Error),

    #[error("failed to convert image {0} to {1}")]
    ConvertImage(PathBuf, PathBuf, #[source] QemuUtilError),
}