view src/xml.rs @ 25:da27437a94b0

Merge remote-tracking branch 'self/master'
author AnaTofuZ <anatofuz@gmail.com>
date Tue, 03 Nov 2020 18:38:03 +0900
parents e8ba0f63c227 0ee235caebc5
children afec42bdd5ab
line wrap: on
line source

use quick_xml::events::{BytesEnd, BytesStart, BytesText, Event};
use quick_xml::{Reader, Writer};
use rand::Rng;
use std::fs;
use std::fs::File;
use std::io::{BufReader, BufWriter, Error};
use std::path::Path;
use uuid::Uuid;

const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
                        abcdefghijklmnopqrstuvwxyz\
                        0123456789)(*^%$#@!~";

const PASSWORD_LEN: usize = 30;

const DOMAIN_XMLNS_QEMU: (&str, &str) =
    ("xmlns:qemu", "http://libvirt.org/schemas/domain/qemu/1.0");

const IE_VIRSH_TEMPLATE_VM_NAME: &[u8; 17] = b"ie-virsh-template";
const VNC_XML_TAG: &[u8; 8] = b"graphics";

const ROOT_START_TAG: &[u8; 6] = b"domain";

const QEMU_COMMAND_LINE_TAG: &[u8; 16] = b"qemu:commandline";
const QEMU_ARG_TAG: &[u8; 8] = b"qemu:arg";

const TEMPLATE_XML_FILE: &str = "/etc/libvirt/template.xml";

const LIBVIRT_XML_DIR: &str = "/etc/libvirt/qemu";
const QCOW2_PATH: &str = "/mnt/ie-virsh";

pub struct GenerateVMArg {
    VMName: String,
    Qcow2PATH: String,
    XMLPATH: String,
    VNCPassword: String,
    IsDebug: bool,
    TCPPort: u64,
}

impl GenerateVMArg {
    pub fn new(user_name: &str, vm_name: &str, is_debug: bool) -> GenerateVMArg {
        let year = user_name.chars().skip(1).take(2).collect::<String>();
        let affilication = if year.parse::<u8>().is_ok() {
            // /etc/libvirt/qemu/e19/e195729
            user_name.chars().take(3).collect::<String>()
        } else {
            "teacher".to_string()
        };

        let xml_dir = format!("{}/{}/{}", LIBVIRT_XML_DIR, affilication, user_name);
        let xml_path = format!("{}/{}.xml", xml_dir.clone(), vm_name);

        if !Path::new(&xml_dir).exists() {
            fs::create_dir_all(xml_dir).ok();
        }

        let qcow2_dir = format!("{}/{}/{}", QCOW2_PATH, affilication, user_name);

        let qcow2_path = format!("{}/{}.qcow2", qcow2_dir.clone(), vm_name);

        if !Path::new(&qcow2_dir).exists() {
            fs::create_dir_all(qcow2_dir).ok();
        }

        let pw = generate_pw();

        GenerateVMArg {
            VMName: vm_name.to_string(),
            Qcow2PATH: qcow2_path,
            XMLPATH: xml_path,
            VNCPassword: pw,
            IsDebug: is_debug,
            TCPPort: 90,
        }
    }

    pub fn generate(self) -> Result<String, Error> {
        let mut reader = Reader::from_reader(BufReader::new(File::open(TEMPLATE_XML_FILE)?));

        print!("generate xml :{}\n", self.XMLPATH);
        let mut writer = Writer::new(BufWriter::new(File::create(self.XMLPATH.clone()).unwrap()));
        let mut buf = Vec::new();
        loop {
            match reader.read_event(&mut buf) {
                Ok(Event::Start(ref e)) if e.name() == b"uuid" => {
                    writer
                        .write_event(Event::Start(e.clone()))
                        .expect("faild write event");
                    reader
                        .read_event(&mut Vec::new())
                        .expect("faild read event");
                    let vm_uuid = Uuid::new_v4().to_string();
                    let elem = BytesText::from_plain_str(&vm_uuid);
                    writer.write_event(Event::Text(elem)).unwrap();
                }

                Ok(Event::Start(ref e)) if (e.name() == ROOT_START_TAG && self.IsDebug) => {
                    let mut elem = e.clone();
                    elem.push_attribute(DOMAIN_XMLNS_QEMU);
                    writer.write_event(Event::Start(elem)).unwrap();

                    let qemu_command_line_start = BytesStart::borrowed_name(QEMU_COMMAND_LINE_TAG);
                    writer
                        .write_event(Event::Start(qemu_command_line_start))
                        .unwrap();

                    for value in ["-S", "-gdb", "tcp"].iter() {
                        let mut qemu_elem = BytesStart::borrowed_name(QEMU_ARG_TAG);
                        let v: &str = &value;
                        qemu_elem.push_attribute(("value", v));
                        writer.write_event(Event::Empty(qemu_elem)).unwrap();
                    }

                    let qemu_command_line_end = BytesEnd::borrowed(QEMU_COMMAND_LINE_TAG);
                    writer
                        .write_event(Event::End(qemu_command_line_end))
                        .unwrap();
                }

                Ok(Event::Empty(ref e)) if e.name() == VNC_XML_TAG => {
                    let mut elem = e.clone();
                    let pw: &str = &self.VNCPassword;
                    elem.push_attribute(("passwd", pw));
                    writer.write_event(Event::Empty(elem)).ok();
                }

                Ok(Event::Empty(ref e)) if (e.name() == b"source") => {
                    let mut elem = e.clone();
                    let is_qcow_file = elem
                        .attributes()
                        .find(|attr| attr.as_ref().unwrap().key == b"file");
                    if is_qcow_file.is_some() {
                        elem.clear_attributes();
                        let qcow2_path: &str = &self.Qcow2PATH;
                        elem.push_attribute(("file", qcow2_path));
                    }
                    writer.write_event(Event::Empty(elem)).ok();
                }

                Ok(Event::Text(ref e)) if e.escaped() == IE_VIRSH_TEMPLATE_VM_NAME => {
                    let elem = BytesText::from_plain_str(&self.VMName);
                    writer.write_event(Event::Text(elem)).unwrap();
                }
                Ok(Event::Eof) => break,
                // you can use either `e` or `&e` if you don't want to move the event
                Ok(e) => assert!(writer.write_event(&e).is_ok()),
                Err(e) => panic!("Error at position {}: {:?}", reader.buffer_position(), e),
            }
            buf.clear();
        }
        print!("generate xml : {}\n", self.XMLPATH);
        print!("vnc password : {}\n", self.VNCPassword);
        Ok(self.XMLPATH)
    }
}

fn generate_pw() -> String {
    let mut rng = rand::thread_rng();

    (0..PASSWORD_LEN)
        .map(|_| {
            let idx = rng.gen_range(0, CHARSET.len());
            CHARSET[idx] as char
        })
        .collect()
}