blctl 0.1.2

// cmdline.rs
//
// Copyright 2019 Alberto Ruiz <aruiz@gnome.org>
//
// This file is free software; you can redistribute it and/or modify it
// under the terms of the GNU Lesser General Public License as
// published by the Free Software Foundation; either version 2.1 of the
// License, or (at your option) any later version.
//
// This file is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
// SPDX-License-Identifier: LGPL-2.1-or-later

use std::io;

pub type CmdlineParam = Vec<Option<String>>;
pub type Cmdline = Vec<(String, CmdlineParam)>;

pub fn cmdline_parse (args: &str) -> std::io::Result<(Cmdline)> {
  match Cmdline::parse(args) {
    Ok(cmdline) => Ok(cmdline),
    Err(_) => {
      Err(io::Error::new(io::ErrorKind::InvalidData,
          format!("could not parse kernelopts variable")))
    }
  }
}

pub trait VecOfTuplesAsDict {
  fn get<'a> (&'a mut self, key: &String) -> Option<&'a mut CmdlineParam>;
  fn take_from_key (&mut self, key :&String) -> Option<CmdlineParam>;
  fn entry_or_insert<'a> (&'a mut self, key: String, insert: CmdlineParam) -> &'a mut CmdlineParam;
}

impl VecOfTuplesAsDict for Cmdline {
  fn get<'a> (&'a mut self, key: &String) -> Option<&'a mut CmdlineParam> {
    for item in self {
      if &item.0 == key {
        return Some(&mut item.1)
      }
    }
    None
  }

  fn take_from_key (&mut self, key: &String) -> Option<CmdlineParam> {
    match { self.iter().position(|x| &x.0 == key) } {
      Some(pos) => Some(self.remove(pos).1),
      None => None
    }
  }

  fn entry_or_insert<'a> (&'a mut self, key: String, insert: CmdlineParam) -> &'a mut CmdlineParam {
    let pos = { self.iter().position(|(k,_)| {k == &key}) };

    match pos {
      Some(index) => {
        &mut self[index].1
      }
      None => {
        self.push((key, insert));
        let len = { self.len() };
        &mut self[len-1].1
      }
    }
  }
}

pub trait CmdlineStore {
  fn cmdline_store(&mut self, cmdline: &Cmdline) -> std::io::Result<()>;
  fn cmdline (&self) -> std::io::Result<Cmdline>;
}

pub trait CmdlineMut {
  fn parse (buffer: &str) -> Result<Cmdline, (usize, &str)>;
  fn render (&self) -> Result<String, &str>;
  fn add_param (&mut self, key: String, value: Option<String>);
}

impl CmdlineMut for Cmdline {
  fn parse (buffer: &str) -> Result<Cmdline, (usize, &str)> {
    #[derive(Debug)]
    enum Scope {
      InValueQuoted,
      InValueUnquoted,
      InKey,
      InEqual,
      InSpace
    };

    let mut key = String::new();
    let mut value = String::new();

    let mut result = Cmdline::new();
    let mut scope = Scope::InSpace;

    let mut i: usize = 0;
    for c in buffer.chars() {
      match c {
        ' ' => {
          match scope {
            Scope::InValueQuoted => { value.push(c); },
            Scope::InValueUnquoted => {
              result.add_param(key.drain(..).collect(), Some(value.drain(..).collect()));
              scope = Scope::InSpace;
            },
            Scope::InSpace => { },
            Scope::InEqual => { return Err((i, "empty parameter value")); },
            Scope::InKey => { result.add_param(key.drain(..).collect(), None); }          }
        },
        '"' => {
          match scope {
            Scope::InValueQuoted => { scope = Scope::InValueUnquoted; }
            Scope::InEqual => { scope = Scope::InValueQuoted; }
            Scope::InKey => { return Err((i, "quote in parameter name")); }
            Scope::InValueUnquoted => { scope = Scope::InValueQuoted; }
            Scope::InSpace => { return Err((i, "quote after unquoted space")); }
          }
        },
        '=' => {
          match scope {
            Scope::InKey => { scope = Scope::InEqual; },
            Scope::InValueQuoted | Scope::InValueUnquoted => { value.push(c); }
            Scope::InEqual => { scope = Scope::InValueUnquoted; value.push(c) }
            Scope::InSpace => { return Err((i, "equals after space")); }
          }
        },
        _ => {
          match scope {
            Scope::InKey => { key.push(c); },
            Scope::InValueQuoted  => { value.push(c); },
            Scope::InValueUnquoted => { value.push(c); },
            Scope::InSpace => { scope = Scope::InKey; key.push(c); }
            Scope::InEqual => { scope = Scope::InValueUnquoted; value.push(c); }
          }
        }
      };
      i += 1;
    }

    match scope {
      Scope::InKey => { result.add_param(key.drain(..).collect(), None); }
      Scope::InValueQuoted => { return Err((i, "unclosed quote in parameter value")); }
      Scope::InValueUnquoted => { result.add_param(key.drain(..).collect(), Some(value.drain(..).collect()))}
      Scope::InEqual => { return Err((i, "empty parameter value")); }
      Scope::InSpace => {}
    }

    Ok(result)
  }

  fn add_param(&mut self, key: String, value: Option<String>) {
    let vec = self.entry_or_insert(key, Vec::new());
    vec.push(value);
  }

  fn render (&self) -> Result<String, &str> {
    let mut render = String::new();
    for (param, values) in self {
      for value in values {
        match value {
          Some(value) => {
            render.push_str(&param);
            render.push('=');
            if value.contains('"') {
              return Err("cannot escape quote character");
            }
            if value.contains(' ') {
              render.push('"');
              render.push_str(&value);
              render.push('"');
            } else {
              render.push_str(&value);
            }

          },
          _ => { render.push_str(&param); }
        }
        render.push(' ');
      }
    }
    render.pop();

    Ok(render)
  }
}

pub trait CmdlineHandler {
  fn cmdline_render(&self) -> std::io::Result<String>;
  fn cmdline_set (&mut self, params: &[String]) -> std::io::Result<()>;
  fn cmdline_get (&self, param: &String) -> std::io::Result<CmdlineParam>;
  fn cmdline_add (&mut self, params: &[String]) -> std::io::Result<()>;
  fn cmdline_remove(&mut self, params: &[String]) -> std::io::Result<()>;
  fn cmdline_clear (&mut self, param: &[String]) -> std::io::Result<()>;
}

impl<T> CmdlineHandler for T
where
  T: CmdlineStore {
  fn cmdline_render(&self) -> std::io::Result<String> {
    let cmdline = self.cmdline()?;
    cmdline.render().or_else(|e| {
      Err(io::Error::new(io::ErrorKind::InvalidData,
                    format!("{}", e)))
    })
  }


  fn cmdline_set (&mut self, params: &[String]) -> std::io::Result<()> {
    let params = cmdline_parse(&params.join(" ").as_str())?;
    let mut cmdline: Cmdline = match self.cmdline() {
      Ok(cmdline) => Ok(cmdline),
      Err(e) => { match e.kind() {
          std::io::ErrorKind::InvalidInput => {
            Ok(Vec::new())
          },
          _ => Err(e)
        }
      }
    }?;
    let mut commit = false;
    for (set_key, set_values) in params {
      let mut values = cmdline.entry_or_insert(set_key, Vec::new());

      for set_value in set_values {
        if ! values.contains(&set_value) {
          commit = true;
          values.pop(); // We replace the last instance
          values.push(set_value);
        }
      }
    }
    match commit {
      true => self.cmdline_store(&cmdline),
      false => Ok(())
    }
  }

  fn cmdline_get (&self, param: &String) -> std::io::Result<CmdlineParam> {
    let mut cmdline = self.cmdline()?;
    match { cmdline.take_from_key(param) } {
      Some(values) => { Ok(values) }
      None => {
        Err(io::Error::new(io::ErrorKind::InvalidData,
                    format!("'{}' parmeter not present in kernelopts bootloader environment variable", param)))
      }
    }
  }

  fn cmdline_add (&mut self, params: &[String]) -> std::io::Result<()> {
    let mut cmdline = self.cmdline()?;
    let params = cmdline_parse(&params.join(" ").as_str())?;
    let mut commit = false;
    for (add_key, add_values) in params {
      let mut found = false;
      for (k,params_for_key) in cmdline.iter_mut() {
        if k != &add_key {
          continue;
        };
        found = true;
        for val in &add_values {
          if ! params_for_key.contains(&val) {
            commit = true;
            params_for_key.push(val.clone());
          }
        }
      };

      if ! found {
        commit = true;
        cmdline.push((add_key.clone(), add_values.clone()));
      }
    }

    match commit {
      true => self.cmdline_store(&cmdline),
      false => Ok(())
    }
  }

  fn cmdline_remove (&mut self, params: &[String]) -> std::io::Result<()> {
    let mut cmdline = self.cmdline()?;
    let params = cmdline_parse(&params.join(" ").as_str())?;
    let mut commit = false;

    for (rem_key, rem_values) in params {
      let mut params_for_key = cmdline.entry_or_insert(rem_key, rem_values.clone());
      for val in rem_values {
        while let Some(index) = params_for_key.iter().position(|v| v == &val) {
          commit = true;
          params_for_key.remove(index);
        }
      }
    }

    match commit {
      true => self.cmdline_store(&cmdline),
      false => Ok(())
    }
  }

  fn cmdline_clear (&mut self, params: &[String]) -> std::io::Result<()> {
    let mut cmdline = self.cmdline()?;
    for param in params {
      if let Some(_) = cmdline.take_from_key(param) {
        if let Err(e) = self.cmdline_store(&cmdline) {
          return Err(e)
        }
      }
    }
    Ok(())
  }
}

#[cfg(test)]
mod bls_tests {
  use std::fs;
  use std::env;
  use bls;
  use blenv;
  use cmdline::Cmdline;
  use cmdline::CmdlineMut;
  use cmdline::CmdlineHandler;

  extern crate tempfile;
  extern crate serial_test_derive;
  use self::serial_test_derive::serial;


  fn tests_init () -> tempfile::TempDir {
    let tmpdir = tempfile::tempdir().expect("Could not create temp dir");
    env::set_var("BLCTL_BLS_DIR", tmpdir.path().as_os_str());

    tmpdir
  }

  fn tests_finalize () {
    env::remove_var("BLCTL_BLS_DIR");
  }

  #[test]
  #[serial]
  fn cmdline_handler_trait_for_bls () {
    let tmpdir = tests_init();

    let mut a_path = tmpdir.path().to_path_buf();
    a_path.push("A.conf");

    // Set arguments
    fs::write(format!("{}/A.conf", tmpdir.path().to_str().unwrap()), "").expect("Could not write test BLS entry");
    let mut a = bls::BLSEntry::new(&String::from("A")).expect("Could not create BLSEntry for A.conf");
    let args = "a=b b c=d".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    a.cmdline_set(&args).expect("Could not set cmdline arguments into BLS entry");

    // Assert with cmdline_get
    assert_eq!(a.cmdline_get(&String::from("a")).expect("Could not get 'a' cmdline argument"),
               vec![Some(String::from("b"))]);
    assert_eq!(a.cmdline_get(&String::from("b")).expect("Could not get 'b' cmdline argument"),
               vec![None]);
    assert_eq!(a.cmdline_get(&String::from("c")).expect("Could not get 'c' cmdline argument"),
               vec![Some(String::from("d"))]);

    // Assert with file contents
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "options a=b b c=d\n");

    // Add arguments
    let args = vec![String::from("a=c"), String::from("a=d"), String::from("a")];
    a.cmdline_add(&args).expect("Could not add 'a=c a=d a' as cmdline args");

    // Assert with cmdline_get
    assert_eq!(a.cmdline_get(&String::from("a")).expect("Could not get 'a' cmdline argument"),
               vec![Some(String::from("b")), Some(String::from("c")), Some(String::from("d")), None]);

    // Assert with cmdline_render
    assert_eq!(a.cmdline_render().expect("Could not render cmdline from env file"),
               String::from("a=b a=c a=d a b c=d"));

    // Assert with file contents
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "options a=b a=c a=d a b c=d\n");

    // Remove one argument assignment
    let args = "a=c".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    a.cmdline_remove(&args).expect("Could not set cmdline arguments into BLS entry");
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "options a=b a=d a b c=d\n");

    // Clear a and b arguments
    let args = "a b".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    a.cmdline_clear(&args).expect("Could not set cmdline arguments into BLS entry");
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "options c=d\n");

    tests_finalize();
  }

  #[test]
  #[serial]
  fn cmdline_handler_trait_for_bls_list () {
    let tmpdir = tests_init();

    let mut a_path = tmpdir.path().to_path_buf();
    let mut b_path = a_path.clone();
    let mut c_path = a_path.clone();
    a_path.push("A.conf");
    b_path.push("B.conf");
    c_path.push("C.conf");
    fs::write(a_path.to_str().expect("Could not turn path into sring"), "").expect("Could not write test BLS entry");
    fs::write(b_path.to_str().expect("Could not turn path into sring"), "").expect("Could not write test BLS entry");
    fs::write(c_path.to_str().expect("Could not turn path into sring"), "").expect("Could not write test BLS entry");

    let mut all = bls::BLSEntry::get_bls_entries().expect("Could not get all BLS entries");
    let args = "a=b b c=d".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    all.cmdline_set(&args).expect("Could not set cmdline arguments into BLS entry");

    // Assert with file contents
    for path in &[&a_path, &b_path, &c_path] {
      assert_eq!(fs::read_to_string(path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "options a=b b c=d\n");
    }

    // Add arguments
    let args = vec![String::from("a=c"), String::from("a=d"), String::from("a")];
    all.cmdline_add(&args).expect("Could not add 'a=c a=d a' as cmdline args");

    // Assert with file contents
    for path in &[&a_path, &b_path, &c_path] {
      assert_eq!(fs::read_to_string(path.as_path()).expect("Could not read BLSEntry file").as_str(),
                 "options a=b a=c a=d a b c=d\n");
    }

    // Remove one argument assignment
    let args = "a=c".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    all.cmdline_remove(&args).expect("Could not set cmdline arguments into BLS entry");
    for path in &[&a_path, &b_path, &c_path] {
      assert_eq!(fs::read_to_string(path.as_path()).expect("Could not read BLSEntry file").as_str(),
                 "options a=b a=d a b c=d\n");
    }

    // Clear a and b arguments
    let args = "a b".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    all.cmdline_clear(&args).expect("Could not set cmdline arguments into BLS entry");
    for path in &[&a_path, &b_path, &c_path] {
      assert_eq!(fs::read_to_string(path.as_path()).expect("Could not read BLSEntry file").as_str(),
                 "options c=d\n");
    }

    tests_finalize();
  }

  #[test]
  #[serial]
  fn cmdline_handler_trait_for_env () {
    let tmpdir = tests_init();

    let mut a_path = tmpdir.path().to_path_buf();
    a_path.push("grubenv");

    let mut env = blenv::Environment::new(None, a_path.to_str().expect("Could not turn env path into sring")).expect("Could not open grubenv file");
    let args = "a=b b c=d".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    env.cmdline_set(&args).expect("Could not set cmdline arguments into BLS entry");

    // Assert with cmdline_get
    assert_eq!(env.cmdline_get(&String::from("a")).expect("Could not get 'a' cmdline argument"),
               vec![Some(String::from("b"))]);
    assert_eq!(env.cmdline_get(&String::from("b")).expect("Could not get 'b' cmdline argument"),
               vec![None]);
    assert_eq!(env.cmdline_get(&String::from("c")).expect("Could not get 'c' cmdline argument"),
               vec![Some(String::from("d"))]);

    // Assert with file contents
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "# GRUB Environment Block\nkernelopts=a=b b c=d\n");

    // Add arguments
    let args = vec![String::from("a=c"), String::from("a=d"), String::from("a")];
    env.cmdline_add(&args).expect("Could not add 'a=c a=d a' as cmdline args");

    // Assert with cmdline_get
    assert_eq!(env.cmdline_get(&String::from("a")).expect("Could not get 'a' cmdline argument"),
               vec![Some(String::from("b")), Some(String::from("c")), Some(String::from("d")), None]);

    // Assert with cmdline_render
    assert_eq!(env.cmdline_render().expect("Could not render cmdline from env file"),
               String::from("a=b a=c a=d a b c=d"));

    // Assert with file contents
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "# GRUB Environment Block\nkernelopts=a=b a=c a=d a b c=d\n");

    // Remove one argument assignment
    let args = "a=c".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    env.cmdline_remove(&args).expect("Could not set cmdline arguments into BLS entry");
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "# GRUB Environment Block\nkernelopts=a=b a=d a b c=d\n");

    // Clear a and b arguments
    let args = "a b".split(" ").map(|arg| { String::from(arg)}).collect::<Vec<String>>();
    env.cmdline_clear(&args).expect("Could not set cmdline arguments into BLS entry");
    assert_eq!(fs::read_to_string(a_path.as_path()).expect("Could not read BLSEntry file").as_str(),
               "# GRUB Environment Block\nkernelopts=c=d\n");

    tests_finalize();
  }

  #[test]
  fn cmdline_parse () {
    let test = "a=test";
    assert_eq!(Cmdline::parse(test), Ok(vec![(String::from("a"),vec![Some(String::from("test"))])]));

    let test = "a=te\"s\"t";
    assert_eq!(Cmdline::parse(test), Ok(vec![(String::from("a"),vec![Some(String::from("test"))])]));

    let test = "a b c";
    assert_eq!(Cmdline::parse(test), Ok(vec![(String::from("a"),vec![None]), (String::from("b"),vec![None]), (String::from("c"),vec![None])]));

    let test = "a=test a a=test2 c a=test3";
    assert_eq!(Cmdline::parse(test), Ok(vec![(String::from("a"),  vec![Some(String::from("test")), None, Some(String::from("test2")), Some(String::from("test3"))]), (String::from("c"), vec![None])]));

    let test = "a=3 =asd";
    assert!(Cmdline::parse(test).is_err());

    let test = "a=3 b= ";
    assert!(Cmdline::parse(test).is_err());

    let test = "a=3 b=";
    assert!(Cmdline::parse(test).is_err());

    let test = "\"quoted param\"=should_error";
    assert!(Cmdline::parse(test).is_err());

    let test = "quot\"ed param\"=should_error";
    assert!(Cmdline::parse(test).is_err());

    let test = "arg1 \"quoted param\"=should_error";
    assert!(Cmdline::parse(test).is_err());

    let test = "param=\"unclosed quote";
    assert!(Cmdline::parse(test).is_err());
  }
}