blctl 0.1.2

// bls.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 cmdline::*;

use std::io;
use std::io::{Read, Write};
use std::path::Path;

use blenv;


#[link(name = "rpm")]
extern {
  fn rpmvercmp(a: *const i8, b: *const i8) -> std::os::raw::c_int;
}

fn compare_kernel_string (a: &std::ffi::CString, b: &std::ffi::CString) -> std::cmp::Ordering {
  match unsafe { rpmvercmp(a.as_ptr(), b.as_ptr()) } {
    1 => { std::cmp::Ordering::Greater },
    0 => { std::cmp::Ordering::Equal },
    -1 => { std::cmp::Ordering::Less },
    _ => { panic!("rpmvercmp() returned an unexpected value"); }
  }
}

#[derive(Debug)]
enum BLSEntryLine {
  Comment(String),
  Command(String,String),
  CommandAndComment(String,String,String),
  Empty
}

type BLSEntryData = Vec<BLSEntryLine>;
type BLSEntryList = Vec<String>;

pub struct BLSEntry {
  bls_dir: String,
  pub name: String
}

impl BLSEntry {

  fn get_full_path(&self) -> String {
    format!("{}/{}", self.bls_dir, self.name)
  }

  fn get_bls_dir() -> String {
    match std::env::var("BLCTL_BLS_DIR") {
      Ok(env_dir) => env_dir,
      _ => { String::from(blenv::BLS_ENTRIES) }
    }
  }

  pub fn new (entry: &String) -> Result<BLSEntry, String> {
    let index = entry.as_str().parse::<usize>();
    let bls_dir = Self::get_bls_dir();

    let entries = match Self::get_bls_entries() {
      Ok(entries) => entries,
      Err(e) => { return Err(format!("could not read bootloader entries from {}: {}", bls_dir, e)); }
    };

    if index.is_ok() {
      let index = index.unwrap();
      if index >= entries.len() {
        return Err(format!("Invalid bootloader index {}, there are only {} entries", index, entries.len()));
      }
      let entry = &entries[index];
      return Ok(BLSEntry {name: entry.clone(), bls_dir: bls_dir});
    }

    // Attempt to identify if it's an BLS entry

    for e in &entries {
      if e == entry {
        let unsubfixed = &entry[0..(entry.len()-5)];
        return Ok(BLSEntry {name: String::from(unsubfixed), bls_dir: bls_dir});
      } else if e == &format!("{}.conf", entry) {
        return Ok(BLSEntry {name: e.clone(), bls_dir: bls_dir});
      }
    }

    // Attempt to parse as path
    let path = std::path::Path::new(&entry);
    if let Ok(abs_path) = path.canonicalize() {
      if abs_path.exists() {
        if abs_path.parent() != Some(Path::new(bls_dir.as_str())) {
          return Err(format!("entry argument {} is not in bootloader entry directory {}", entry, bls_dir));
        }
        if !abs_path.ends_with(".conf") {
          return Err(format!("entry argument {:?} does not have .conf suffix", abs_path));
        }
        if let Some(file_stem) = abs_path.file_stem() {
          if let Some(file_stem) = file_stem.to_str() {
            return Ok(BLSEntry {name: String::from(file_stem), bls_dir: bls_dir});
          }
        }
      }
    }

    // Attempt to parse as kernel entry
    let mut candidates: Vec<String> = Vec::new();
    for e in entries {
      match std::fs::File::open(&format!("{}/{}", bls_dir, e)) {
        Ok(file) => {
          let mut content = String::new();
          let mut reader =  std::io::BufReader::new(file);
          let _ = reader.read_to_string(&mut content);
          for line in content.lines() {
            let line = line.split("#").collect::<Vec<_>>()[0].trim();
            if line.starts_with("version ") && line.ends_with(entry) {
              candidates.push(String::from(e.as_str()));
            }
          }
        }
        _ => {}
      };
    };

    if candidates.len() == 1 {
      return Ok(BLSEntry {name: entry.clone(), bls_dir: bls_dir});
    }

    if candidates.len() > 2 {
      return Err(format!("ERROR: kernel version {} was found in multiple boot entries: {}", entry, candidates.join(" ")));
    }

    Err(format!("ERROR: invalid entry argument {}", entry))
  }

  pub fn get_bls_entries () -> std::io::Result<BLSEntryList> {
    let bls_dir = Self::get_bls_dir();
    let dir = std::fs::read_dir(std::path::Path::new(bls_dir.as_str()))?;
    let mut entries: Vec<std::ffi::CString> = dir.filter(|entry| {entry.is_ok()})
                             .map(|entry| {entry.unwrap().file_name().into_string()})
                             .filter(|entry| {entry.is_ok()})
                             .map(|entry| {entry.unwrap()})
                             .filter(|entry| {entry.ends_with(".conf")})
                             .filter_map(|entry| {
                               if let Ok(entry) = std::ffi::CString::new(entry) {
                                 Some(entry)
                               } else {
                                 None
                               }
                             })
                             .collect();
    // Reverse ordering
    entries.as_mut_slice().sort_by(|a,b| {compare_kernel_string(b,a)});

    let entries = entries.drain(..)
      .map(|c_entry| { c_entry.into_string() })
      .filter_map(|entry| {
        if let Ok(entry) = entry {
          Some(entry)
        } else {
          None
        }
      }).collect();
    Ok(entries)
  }

  fn parse (&self) -> std::io::Result<BLSEntryData> {
    let mut result = Vec::new();
    let mut content = String::new();
    let path = self.get_full_path();
    let mut file = std::fs::File::open(&path)?;
    file.read_to_string(&mut content)?;

    for line in content.lines() {
      let mut comment = None;
      let line = if line.contains("#") {
        let split: Vec<_> = line.splitn(2, "#").collect();
        comment = Some(String::from(split[1]));
        split[0]
      } else {
        line
      };
      if ! line.trim().contains(" ") {
        match comment {
          Some(comment) => { result.push(BLSEntryLine::Comment(comment)); }
          None => { result.push(BLSEntryLine::Empty); }
        }
      } else {
        let mut line: Vec<&str> = line.trim().splitn(2, " ").collect();
        let command = (String::from(line[0]), String::from(line[1]));
        match comment {
          Some(comment) => {
            result.push(BLSEntryLine::CommandAndComment(command.0, command.1, comment));
          }
          None => { result.push(BLSEntryLine::Command(command.0, command.1)); }
        }
      }
    }

    Ok(result)
  }

  fn commit (&self, data: &BLSEntryData) -> std::io::Result<()> {
    let mut content = String::new();
    let path = self.get_full_path();
    let mut file = std::fs::File::create(&path)?;
    for line in data {
      match line {
        BLSEntryLine::Command(command,args) => {
          content.push_str(command);
          content.push_str(" ");
          content.push_str(args);
          content.push_str("\n");
        },
        BLSEntryLine::Comment(comment) => {
          content.push_str("#");
          content.push_str(comment);
          content.push_str("\n");
        },
        BLSEntryLine::CommandAndComment(command,args,comment) => {
          content.push_str(command);
          content.push_str(" ");
          content.push_str(args);
          content.push_str(" #");
          content.push_str(comment);
          content.push_str("\n");
        }
        _ => { content.push_str("\n"); }
      }
    }
    // TODO: Write and then move
    file.write_all(content.as_bytes())
  }

  pub fn get(&self, command: &str) -> std::io::Result<String> {
    let data = self.parse()?;
    let matches: Vec<_> = data.iter()
      .filter(|x| match x {
        BLSEntryLine::CommandAndComment(c,_,_) |
        BLSEntryLine::Command(c,_) => { c.as_str() == command }
        _ => { false }
      }).map(|x| match x {
        BLSEntryLine::CommandAndComment(_,arg,_) |
        BLSEntryLine::Command(_,arg) => {
          arg
        }
        _ => { panic!("unreachable code"); }
      })
      .collect();

    match matches.len() {
      0 => { Err(io::Error::new(io::ErrorKind::InvalidInput,
                    format!("'{}' does not contain a {} command", self.get_full_path(), command))) },
      _ => { Ok(matches[0].clone()) }
    }
  }

  pub fn set(&self, command: &str, value: String) -> std::io::Result<()> {
    let mut data = self.parse()?;
    let command = String::from(command);
    let pos = data.iter().position(|line| match line {
        BLSEntryLine::Command(k,_) | BLSEntryLine::CommandAndComment(k,_,_) => {
          k.as_str() == command.as_str()
        }
        _ => false
      });

    if let Some(i) = pos {
      let candidate = match &data[i] {
        BLSEntryLine::Command(_,v) => {
          if v != &value {
              Some(BLSEntryLine::Command(command, value))
            } else {
              None
            }
        }
        BLSEntryLine::CommandAndComment(_,v,comment) => {
          if v != &value {
            Some(BLSEntryLine::CommandAndComment(command, value, comment.clone()))
          } else {
            None
          }
        }
        _ => { None }
      };

      if let Some(line) = candidate {
        data.remove(i);
        data.push(line);
        self.commit(&data)?;
      }
    } else {
      data.push(BLSEntryLine::Command(command, value));
      self.commit(&data)?;
    };
    Ok(())
  }

  pub fn remove(&self, command: &str) -> std::io::Result<()> {
    let mut data = self.parse()?;

    let mut i: usize = 0;
    let mut candidates = Vec::new();
    let mut commit = false;

    for mut line in data.iter() {
      match line {
        BLSEntryLine::Command(c,_) => {
          if c == command {
            candidates.push(i);
          };
        },
        BLSEntryLine::CommandAndComment(c,_,_) => {
          if c == command {
            candidates.push(i);
          };
        },
        _ => {}
      };
      i += 1;
    };

    for i in candidates {
      if ! commit { commit = true };
      match data.remove(i) {
        BLSEntryLine::CommandAndComment(_,_,cmnt) => {
          data.insert(i, BLSEntryLine::Comment(cmnt));
        }
        _ => {}
      };
    };

    if commit {
      self.commit(&data)
    } else {
      Ok(())
    }
  }
}

impl CmdlineStore for BLSEntry {
  fn cmdline_store(&mut self, cmdline: &Cmdline) -> std::io::Result<()> {
    let cmdline_string = match Cmdline::render(cmdline) {
      Ok(cmdline_string) => cmdline_string,
      Err(error) => {
        return Err(io::Error::new(io::ErrorKind::InvalidData,
                    format!("could not store kernel cmdline in '{}': {}", self.name, error)));
      }
    };

    self.set("options", cmdline_string)
  }
  fn cmdline (&self) -> std::io::Result<Cmdline> {
    match Cmdline::parse(self.get("options")?.as_str()) {
      Ok(c) => Ok(c),
      Err(e) => {
        Err(io::Error::new(io::ErrorKind::InvalidData,
                    format!("could not read kernel cmdline from '{}': {}", self.name, e.1)))
      }
    }
  }
}

fn entry_batch_run (entries: &Vec<String>,
                    params: &[String],
                    func: &Fn (&mut BLSEntry, &[String]) -> std::io::Result<()>) -> std::io::Result<()> {

  let errors = entries.iter()
    .map(|e| {
      let mut entry = BLSEntry{name: e.clone(), bls_dir: BLSEntry::get_bls_dir()};
      (func(&mut entry, params), e)
    })
    .filter(|(result, _)| { result.is_err() })
    .map(|(result, entry)| { format!("there was an error trying to modify cmdline parameters in {}: {}",
                                       entry, result.err().unwrap()) })
    .collect::<Vec<String>>()
    .join("\n");

  if errors.len() > 0 {
    Err(io::Error::new(io::ErrorKind::Interrupted,
                       errors))
  } else {
    Ok(())
  }
}

impl CmdlineHandler for Vec<String> {
  fn cmdline_render(&self) -> std::io::Result<String> {
    panic!("not implemented");
  }
  fn cmdline_get (&self, _: &String) -> std::io::Result<CmdlineParam> {
    panic!("not implemented");
  }
  fn cmdline_set (&mut self, params: &[String]) -> std::io::Result<()> {
    entry_batch_run(self, params, &BLSEntry::cmdline_set)
  }
  fn cmdline_add (&mut self, params: &[String]) -> std::io::Result<()> {
    entry_batch_run(self, params, &BLSEntry::cmdline_add)
  }
  fn cmdline_remove(&mut self, params: &[String]) -> std::io::Result<()> {
    entry_batch_run(self, params, &BLSEntry::cmdline_remove)
  }
  fn cmdline_clear (&mut self, params: &[String]) -> std::io::Result<()> {
    entry_batch_run(self, params, &BLSEntry::cmdline_clear)
  }
}

#[cfg(test)]
mod bls_tests {
  use std::fs;
  use std::env;
  use bls;

  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 all_entries() {
    let tmpdir = tests_init();

    let all = bls::BLSEntry::get_bls_entries().expect("Could not list all entries");
    assert_eq!(all.len(), 0);

    fs::write(format!("{}/A.conf", tmpdir.path().to_str().unwrap()), "").expect("Could not write test BLS entry");
    let all = bls::BLSEntry::get_bls_entries().expect("Could not list all entries");
    assert_eq!(all, vec!["A.conf"]);

    fs::write(format!("{}/B.conf", tmpdir.path().to_str().unwrap()), "").expect("Could not write test BLS entry");
    fs::write(format!("{}/C.conf", tmpdir.path().to_str().unwrap()), "").expect("Could not write test BLS entry");
    let all = bls::BLSEntry::get_bls_entries().expect("Could not list all entries");
    assert_eq!(all, vec!["C.conf", "B.conf", "A.conf"]);

    fs::write(format!("{}/C.conf.false", tmpdir.path().to_str().unwrap()), "").expect("Could not write test BLS entry");
    let all = bls::BLSEntry::get_bls_entries().expect("Could not list all entries");
    assert_eq!(all, vec!["C.conf", "B.conf", "A.conf"]);

    tests_finalize();
  }

  #[test]
  #[serial]
  fn new () {
    let tmpdir = tests_init();
    fs::write(format!("{}/A.conf", tmpdir.path().to_str().unwrap()), "").expect("Could not write test BLS entry");
    let a = bls::BLSEntry::new(&String::from("A.conf"));
    assert!(a.is_ok());
    let b = bls::BLSEntry::new(&String::from("B.conf"));
    assert!(b.is_err());

    let zero = bls::BLSEntry::new(&String::from("0")).expect("Get BLSEntry by index");
    assert_eq!(zero.name, String::from("A.conf"));
    let one = bls::BLSEntry::new(&String::from("1"));
    assert!(one.is_err());

    tests_finalize();
  }

  #[test]
  #[serial]
  fn get () {
    let tmpdir = tests_init();
    fs::write(format!("{}/A.conf", tmpdir.path().to_str().unwrap()), "kernel a=1").expect("Could not write test BLS entry");
    let a = bls::BLSEntry::new(&String::from("A")).expect("Could not create BLSEntry instance for existing BLS entry");
    let cmdline = a.get("kernel").expect("Could not get argument for kernel command in BLS entry");
    assert_eq!(cmdline, String::from("a=1"));

    tests_finalize();
  }

  #[test]
  #[serial]
  fn set () {
    let tmpdir = tests_init();
    fs::write(format!("{}/A.conf", tmpdir.path().to_str().unwrap()), "").expect("Could not write test BLS entry");
    let a = bls::BLSEntry::new(&String::from("A")).expect("Could not create BLSEntry instance for existing BLS entry");
    let _ = a.set("kernel", String::from("a=1")).expect("Could not set argument for kernel command in BLS entry");
    let cmdline = a.get("kernel").expect("Could not get argument for kernel command in BLS entry");
    assert_eq!(cmdline, String::from("a=1"));

    tests_finalize();
  }

  #[test]
  #[serial]
  fn remove () {
    let tmpdir = tests_init();
    fs::write(format!("{}/A.conf", tmpdir.path().to_str().unwrap()), "kernel a=1").expect("Could not write test BLS entry");
    let a = bls::BLSEntry::new(&String::from("A")).expect("Could not create BLSEntry instance for existing BLS entry");
    let _ = a.remove("kernel").expect("Could not remove argument for kernel command in BLS entry");
    assert!(a.get("kernel").is_err());

    tests_finalize();
  }
}