Fluentbit config parser

config_parser.rs

use std::fs::File;
use std::io::{self, prelude::*, BufReader};

#[derive(Clone)]
pub struct Property {
    key: String,
    value: String,
}

impl Property {
    fn new(key: String, value: String) -> Property {
        Property { key, value }
    }
}

#[derive(Clone)]
pub struct Section {
    name: String,
    properties: Vec<Property>,
}

fn custom_strip(s: &str) -> &str {
    s.trim_matches(|c: char| c == ' ' || c == '\t')
}

impl Section {
    pub fn new(name: String) -> Section {
        Section {
            name,
            properties: Vec::new(),
        }
    }

    pub fn add_property(&mut self, key: String, value: String) {
        let property = Property::new(key, value);
        self.properties.push(property);
    }

    fn has_properties(&self, properties: &[(String, String)]) -> bool {
        properties.iter().all(|(search_key, search_value)| {
            self.properties
                .iter()
                .any(|property| &property.key == search_key && &property.value == search_value)
        })
    }

    fn get_properties(&self) -> String {
        if self.properties.is_empty() {
            return String::new();
        }

        let max_key_length = self
            .properties
            .iter()
            .map(|prop| prop.key.len())
            .max()
            .unwrap_or(0);
        let indent_length = 4;
        let indent = " ".repeat(indent_length);

        self.properties
            .iter()
            .map(|prop| {
                format!(
                    "{}{}{}  {}\n",
                    indent,
                    prop.key,
                    " ".repeat(max_key_length - prop.key.len()),
                    prop.value
                )
            })
            .collect()
    }

    fn get_section(&self) -> String {
        format!("[{}]\n{}", self.name, self.get_properties())
    }
}

pub struct Parser {
    current_section: Option<Section>,
    sections: Vec<Section>,
    file_path: String,
}

impl Parser {
    pub fn new(file_path: String) -> Parser {
        Parser {
            current_section: None,
            sections: Vec::new(),
            file_path,
        }
    }

    pub fn parse(&mut self) -> io::Result<()> {
        if !std::path::Path::new(&self.file_path).exists() {
            return Err(io::Error::new(io::ErrorKind::NotFound, "File not found"));
        }

        let file = File::open(&self.file_path)?;
        let reader = BufReader::new(file);

        for line in reader.lines() {
            let line = line?;
            let trimmed_line = custom_strip(&line);
            if trimmed_line.is_empty() {
                continue;
            }

            if trimmed_line.starts_with('[') {
                self.parse_section(&trimmed_line);
            } else {
                self.parse_property(&trimmed_line);
            }
        }

        Ok(())
    }

    pub fn section_exists(&self, properties: &[(String, String)]) -> bool {
        self.sections
            .iter()
            .any(|section| section.has_properties(properties))
    }

    pub fn add_section(&mut self, sec: Section) {
        self.sections.push(sec);
    }

    pub fn remove_section(&mut self, properties: &[(String, String)]) {
        self.sections
            .retain(|section| !section.has_properties(properties));
    }

    fn parse_section(&mut self, line: &str) {
        if let Some(section) = self.current_section.take() {
            self.sections.push(section);
        }

        let section_name = line.trim_matches(|c| c == '[' || c == ']' || c == '\n');
        self.current_section = Some(Section::new(section_name.to_string()));
    }

    fn parse_property(&mut self, line: &str) {
        if line.starts_with('#') || self.current_section.is_none() {
            return;
        }

        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() < 2 {
            return;
        }

        let key = parts[0].to_string();
        let value = parts[1..].join(" ");
        self.current_section
            .as_mut()
            .unwrap()
            .add_property(key, value);
    }

    pub fn save(&self) -> io::Result<()> {
        let mut file = File::create(format!("{}", self.file_path))?;
        for section in &self.sections {
            writeln!(file, "{}", section.get_section())?;
        }

        if let Some(current_section) = &self.current_section {
            writeln!(file, "{}", current_section.get_section())?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    fn setup_with_content(contents: &str) -> NamedTempFile {
        let mut file = NamedTempFile::new().expect("Failed to create temp file");
        writeln!(file, "{}", contents).expect("Failed to write to temp file");
        file.flush().expect("Failed to flush temp file");

        file
    }

    #[test]
    fn test_parse() {
        let file = setup_with_content("[SECTION]\n    Foo         Bar\n[]");
        let file_path = file.path().to_str().unwrap().to_string();
        let mut parser = Parser::new(file_path);
        parser.parse().expect("Failed to parse");

        assert_eq!(parser.sections.len(), 1);
    }

    #[test]
    fn test_section_exists() {
        let file = setup_with_content("[SECTION]\n    Foo         Bar\n[]");
        let file_path = file.path().to_str().unwrap().to_string();
        let mut parser = Parser::new(file_path);
        parser.parse().expect("Failed to parse");

        assert!(parser.section_exists(&[("Foo".to_string(), "Bar".to_string())]));
    }

    #[test]
    fn test_add_section_and_save() {
        let file = setup_with_content("[SECTION]\n    Foo         Bar\n[]");
        let file_path = file.path().to_str().unwrap().to_string();

        let mut parser = Parser::new(file_path.clone());
        let mut section = Section::new("NEW_SECTION".to_string());
        section.add_property("new_key".to_string(), "new_value".to_string());
        parser.add_section(section);
        parser.save().expect("Failed to save");

        let saved_contents = std::fs::read_to_string(file_path).expect("Failed to read saved file");
        assert!(saved_contents.contains("[NEW_SECTION]"));
        assert!(saved_contents.contains("new_key  new_value"));
    }

    #[test]
    fn test_remove_section_and_save() {
        let file = setup_with_content("[SECTION]\n    Foo         Bar\n[]");
        let file_path = file.path().to_str().unwrap().to_string();

        let mut parser = Parser::new(file_path.clone());
        parser.remove_section(&[("Foo".to_string(), "Bar".to_string())]);
        parser.save().expect("Failed to save");

        let saved_contents = std::fs::read_to_string(file_path).expect("Failed to read saved file");
        assert!(!saved_contents.contains("[SECTION]"));
    }
}