From 32612f0308b636d238625a52ad449ed5ec5fb5be Mon Sep 17 00:00:00 2001
From: Thang Nguyen <thangnguyenuvm@gmail.com>
Date: Wed, 3 Apr 2024 07:36:27 -0400
Subject: [PATCH] Added flag to toggle groups

---
 src/lib.rs  | 28 ++++++++++++++++++++++++++-
 src/main.rs | 55 +++++++++++++++++++++++++++++++++++++----------------
 2 files changed, 66 insertions(+), 17 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 51e4ad5..db4376d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -21,6 +21,17 @@ pub type HostnameGroups = Vec<HostnameGroup>;
 pub struct HostnameGroup {
     pub name: String,
     pub hostnames: Vec<String>,
+    pub disabled: Option<bool>,
+}
+
+impl HostnameGroup {
+    pub fn new(name: String, hostnames: Vec<String>) -> Self {
+        Self {
+            name,
+            hostnames,
+            disabled: Some(false),
+        }
+    }
 }
 
 impl fmt::Display for HostnameGroup {
@@ -30,7 +41,12 @@ impl fmt::Display for HostnameGroup {
             .iter()
             .map(|hostname| "  ".to_string() + hostname)
             .collect();
-        write!(f, "{}\n{}\n", self.name, padded_hostnames.join("\n"))
+        let disabled = if self.disabled.unwrap_or(false) {
+            " (disabled)"
+        } else {
+            ""
+        };
+        write!(f, "{}{}\n{}\n", self.name, disabled, padded_hostnames.join("\n"))
     }
 }
 
@@ -44,6 +60,7 @@ impl std::default::Default for HostnameGroup {
                 "twitter.com".to_string(),
                 "linkedin.com".to_string(),
             ],
+            disabled: Some(false),
         }
     }
 }
@@ -98,6 +115,9 @@ pub fn construct_refocus_line(hostgroups: &HostnameGroups) -> String {
     info!("Generating Refocus hosts line");
     let mut refocus_line = "127.0.0.1 ".to_owned() + HOSTNAME_ANCHOR + " ";
     for group in hostgroups {
+        if group.disabled.unwrap_or(false) {
+            continue;
+        }
         for hostname in &group.hostnames {
             refocus_line.push_str(hostname);
             refocus_line.push(' ');
@@ -160,3 +180,9 @@ pub fn overwrite_config_file(hostgroups: &HostnameGroups) -> Result<(), Box<dyn
     fs::write(path, content)?;
     Ok(())
 }
+
+pub fn split_args(args: &str) -> Vec<String> {
+    args.split(',')
+        .map(|x| x.to_lowercase())
+        .collect()
+}
diff --git a/src/main.rs b/src/main.rs
index e95a88b..965c592 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -29,6 +29,10 @@ struct Args {
     #[arg(short, long, default_value_t = String::from(""))]
     delete: String,
 
+    /// toogle hostgroup(s)
+    #[arg(short, long, default_value_t = String::from(""))]
+    toggle: String,
+
     /// execute refocus
     #[arg(short, long, default_value_t = true)]
     execute: bool,
@@ -73,16 +77,43 @@ fn main() {
         process::exit(1);
     }
 
+    if !args.toggle.is_empty() {
+        match read_hostname_groups_config() {
+            Ok(mut groups) => {
+                let toggle_groups = split_args(&args.toggle);
+
+                for toggle_group in &toggle_groups {
+                    for group in groups.iter_mut() {
+                        if group.name.to_lowercase() == toggle_group.to_lowercase() {
+                            group.disabled = Some(!group.disabled.unwrap_or(false));
+                            println!("Toggled group: {}. Disabled: {}", group.name, group.disabled.unwrap_or(false));
+                        }
+                    }
+                }
+
+                if overwrite_config_file(&groups).is_err() {
+                    eprintln!("Failed to update hostname groups config");
+                    process::exit(1);
+                }
+            }
+            Err(e) => {
+                eprintln!("Failed to read hostname groups config: {}", e);
+                process::exit(1);
+            }
+        }
+    }
+
     if !args.group.is_empty() && !args.add.is_empty() {
         match read_hostname_groups_config() {
             Ok(mut groups) => {
-                let new_hostnames: Vec<String> = args
-                    .add
-                    .replace(' ', "")
-                    .split(',')
+                let new_hostnames = split_args(
+                        &args.add
+                        .trim()
+                    )
+                    .into_iter()
+                    // .replace(' ', "")
                     .filter(|hostname| hostname.contains('.'))
-                    .map(|hostname| hostname.to_lowercase())
-                    .collect();
+                    .collect::<Vec<String>>();
 
                 if let Some(group) = groups
                     .iter_mut()
@@ -90,10 +121,7 @@ fn main() {
                 {
                     group.hostnames.extend(new_hostnames);
                 } else {
-                    groups.push(HostnameGroup {
-                        name: args.group,
-                        hostnames: new_hostnames,
-                    })
+                    groups.push(HostnameGroup::new(args.group, new_hostnames));
                 }
 
                 if overwrite_config_file(&groups).is_err() {
@@ -111,12 +139,7 @@ fn main() {
     if !args.delete.is_empty() {
         match read_hostname_groups_config() {
             Ok(mut groups) => {
-                let delete_hostnames: Vec<String> = args
-                    .delete
-                    .replace(' ', "")
-                    .split(',')
-                    .map(|hostname| hostname.to_lowercase())
-                    .collect();
+                let delete_hostnames: Vec<String> = split_args(&args.delete);
 
                 for delete_hostname in delete_hostnames {
                     for group in groups.iter_mut() {