path: root/src/star.rs

/*!
  * `Star`s are where [crate::planet::Planet]s orbit (the physical hypervisors that
  * libvirtd connects to
  */
use crate::planet::*;
use crate::errors::Error;
use crate::ship::Ship;
use serde::{Serialize, Deserialize};
use virt::{connect::Connect, domain::Domain};
use std::process::{ExitStatus};
use std::os::unix::process::ExitStatusExt;
use rand::Rng;

use tokio::{process::Command};

#[derive(Serialize, Deserialize, Debug)]
pub enum Address {
    IP(String),
    Domain(String)
}

impl ToString for Address {
    fn to_string(&self) -> String {
        match self {
            Address::IP(s) => s.clone(),
            Address::Domain(s) => s.clone(),
        }
    }
}

/// Defines a "star" where [crate::planet::Planet]s orbit
#[derive(Debug, Serialize, Deserialize)]
pub struct Star {
    /// Hostname
    pub name: String,
    /// FQDN or IP address, a way to talk to the planet
    pub address: String,

    /// Whether or not the Planet is local (same machine) or remote (networked machine)
    pub remote: bool,

    /// Connection to the House, if available
    #[serde(skip)]
    con: Option<Connect>,
}

impl Star { 
    /// Creates a new Planet based on a libvirtd connect URL
    ///
    /// Example:
    /// ```
    /// use waifulib::planet::Planet;
    /// let mut h = Planet::new("test:///default".to_string()).unwrap();
    /// ```
    pub fn new(url: String) -> Result<Self, Error> {
        let c = Connect::open(&url.clone())?;

        let remote = if url.contains("qemu:///") || url.contains("localhost") || url.contains("127.0.0.1") {
            true
        } else {
            false
        };

        // If the connection succeeds, we've got one!
        Ok(Self {
            name: c.get_hostname()?,
            address: c.get_uri()?,
            remote,
            con: Some(c),
        })
    }

    /// Lists the "inhabitants" of the House (the [Waifu]s on the machine
    ///
    /// ```
    /// use crate::star::Star;
    /// let mut h = Star::new("test:///default".to_string()).unwrap();
    ///
    /// assert_eq!(h.inhabitants().unwrap().len(), 1);
    /// ```
    pub fn inhabitants(&mut self) -> Result<Vec<Planet>, Error> {
        match &self.con {
            Some(c) => {
                let domains = c.list_all_domains(0)?;

                let mut planets: Vec<Planet> = Vec::new();

                for d in domains.iter() {
                    planets.push(d.clone().try_into()?);
                }
                Ok(planets)
            },
            None => {
                return Err(Error::Connection("Domain connection was None".to_string()));
            }
        }
    }

    pub fn find_planet(&mut self, uuid: String) -> Result<Planet, Error> {
        let inhab = self.inhabitants()?;

        for i in inhab {
            if i.uuid == uuid {
                return Ok(i);
            }
        }
        Err(Error::Other(String::from("Not found")))
    }

    /// Creates a new Planet orbiting the Star, taking care of everything needed to make the VM run
    /// fine
    ///
    /// If the installation image doesn't exist in the default libvirtd pool, this will fail with
    /// [`Error::MissingImage`][crate::errors::Error::MissingImage].
    ///
    pub async fn planet(&mut self, name: String, max_mem: Memory, max_cpus: CpuCount, disk_size_mb: u64, ship: Ship) -> Result<Planet, Error> {
        // Check for image on host machine
        
        if self.remote {
            let mut output = Command::new("ssh")
                .args([
                      "-oStrictHostKeyChecking=accept-new",
                      &self.address.clone(),
                      "stat",
                      &format!("/var/lib/libvirt/images/{}", ship.make_pretty_name().clone())
                ])
                .output()
                .await?;

            if output.status != ExitStatus::from_raw(0) {
               return Err(Error::MissingImage(ship.name.clone()));
            }

            // Allocate VM disk
            output = Command::new("ssh")
                .args([
                      "-oStrictHostKeyChecking=accept-new",
                      &self.address.clone(),
                      "qemu-img",
                      "create",
                      "-f",
                      "qcow2",
                      &format!("/var/lib/libvirt/images/{}.qcow2", name.clone()),
                      &format!("{}M", disk_size_mb)
                ])
                .output()
                .await?;

            if output.status != ExitStatus::from_raw(0) {
                return Err(Error::Allocation(String::from_utf8(output.stdout).unwrap()));
            }
        } else {
            // It's local
            let mut output = Command::new("stat")
                .args([
                      &format!("/var/lib/libvirt/images/{}", ship.make_pretty_name().clone()),
                ])
                .output()
                .await?;

            if output.status != ExitStatus::from_raw(0) {
                return Err(Error::MissingImage(ship.name.clone()));
            }

            output = Command::new("qemu-img")
                .args([
                      "create",
                      "-f",
                      "qcow2",
                      &format!("/var/lib/libvirt/images/{}.qcow2", name.clone()),
                      &format!("{}M", disk_size_mb)
                ])
                .output()
                .await?;

            if output.status != ExitStatus::from_raw(0) {
                return Err(Error::Allocation(String::from_utf8(output.stdout).unwrap()));
            }
        }

        let uuid = uuid::Uuid::new_v4();
        
        // Let's get that XML ready
        let mut buf: Vec<u8> = vec![];

        crate::templates::vm_xml(
            &mut buf,
            name.clone(),
            uuid.to_string(),
            random_mac().clone(),
            true,
            max_mem.0,
            max_cpus.0,
            "blank".to_string()
            )?;

        let buf = String::from_utf8(buf).unwrap();

        println!("{}", buf);

        let dom: Domain = match &self.con {
                Some(c) => {
                    let dom = Domain::define_xml(&c, &buf)?;

                    dom.create()?;

                    dom
                },
                None => {
                    return Err(Error::Connection("Connection was None".to_string()));
                }
            };

        dom.try_into()
    }

}

fn random_mac() -> String {
    let mut addr = rand::thread_rng().gen::<[u8; 6]>();
    addr[0] = (addr[0] | 2) & 0xfe;
    mac_address::MacAddress::new(addr).to_string()
}

#[cfg(test)]
mod test {
    use super::*;
    /// Kind of a stupid test, but just a sanity check to make sure that [ToString] impl up above
    /// is still working
    #[test]
    fn addr_to_string() {
        let a: Address = Address::IP("127.0.0.1".to_string());
        let d: Address = Address::Domain("example.com".to_string());

        assert_eq!(String::from("127.0.0.1"), a.to_string());
        assert_eq!("example.com".to_string(), d.to_string());
    }

    #[test]
    fn connect_to_house() {
        let _: Star = Star::new("test:///default".to_string()).unwrap();
    }

    #[test]
    fn list_inhabitants() {
        let mut h: Star = Star::new("test:///default".to_string()).unwrap();

        assert_eq!(h.inhabitants().unwrap().len(), 1);        
    }

    #[test]
    fn get_planet() {
        let mut s: Star = Star::new("test:///default".to_string()).unwrap();

        let t = s.inhabitants().unwrap()[0].clone();

        assert_eq!(s.find_planet(t.uuid.clone()).unwrap(), t);
    }

}