From 96f1f8ae811f58bd58910f3e01c44cc3ebbfa063 Mon Sep 17 00:00:00 2001
From: Peter Morrow <peter@graphiant.com>
Date: Thu, 13 Jun 2024 14:11:39 +0100
Subject: [PATCH] priorities: support filtering on log priority

---
 src/main.rs | 119 +++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 100 insertions(+), 19 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index cf21464..ec71261 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -38,6 +38,22 @@ struct Opt {
     /// Filter messages based on this regex pattern.
     #[structopt(short, long)]
     pattern: Option<String>,
+
+    /// Filter messages at this priority.
+    #[structopt(short, long)]
+    priority: Option<String>,
+}
+
+#[derive(Eq, PartialEq, Hash, Clone, Debug, Ord, PartialOrd)]
+enum Priority {
+    EMERGENCY,
+    ALERT,
+    CRITICAL,
+    ERROR,
+    WARN,
+    NOTICE,
+    INFO,
+    DEBUG,
 }
 
 #[derive(Eq, PartialEq, Hash, Clone)]
@@ -47,7 +63,7 @@ struct Message {
     /// The process that generated the message.
     process: String,
     /// Priority the message was sent at.
-    priority: String,
+    priority: Priority,
 }
 
 struct JournalStat {
@@ -69,6 +85,8 @@ struct JournalStat {
     total_msgs: u64,
     // Regex to match on.
     regex: Option<Regex>,
+    // Priority to match on.
+    priority: Option<Priority>,
 }
 
 #[derive(Tabled)]
@@ -77,7 +95,7 @@ struct TopTalkerTableEntry<'a> {
     Rank: usize,
     Frequency: u32,
     Process: &'a str,
-    Priority: String,
+    Priority: &'a str,
     Message: &'a str,
 }
 
@@ -97,6 +115,53 @@ struct SizeTableEntry<'a> {
     Message: &'a str,
 }
 
+impl From<u8> for Priority {
+    fn from(p: u8) -> Self {
+        match p {
+            0 => Self::EMERGENCY,
+            1 => Self::ALERT,
+            2 => Self::CRITICAL,
+            3 => Self::ERROR,
+            4 => Self::WARN,
+            5 => Self::NOTICE,
+            6 => Self::INFO,
+            7 => Self::DEBUG,
+            _ => todo!(),
+        }
+    }
+}
+
+impl From<&str> for Priority {
+    fn from(p: &str) -> Self {
+        match p {
+            "0" | "emergency" => Self::EMERGENCY,
+            "1" | "alert" => Self::ALERT,
+            "2" | "critical" => Self::CRITICAL,
+            "3" | "error" => Self::ERROR,
+            "4" | "warn" => Self::WARN,
+            "5" | "notice" => Self::NOTICE,
+            "6" | "info" => Self::INFO,
+            "7" | "debug" => Self::DEBUG,
+            _ => todo!(),
+        }
+    }
+}
+
+impl Priority {
+    fn as_str(&self) -> &'static str {
+        match self {
+            Self::EMERGENCY => "emergency",
+            Self::ALERT => "alert",
+            Self::CRITICAL => "critical",
+            Self::ERROR => "error",
+            Self::WARN => "warn",
+            Self::NOTICE => "notice",
+            Self::INFO => "info",
+            Self::DEBUG => "debug",
+        }
+    }
+}
+
 impl JournalStat {
     /// Create a new JournalStat struct.
     pub fn new(path: &Path) -> Result<Self, systemd::Error> {
@@ -117,6 +182,7 @@ impl JournalStat {
             per_process: HashMap::new(),
             total_msgs: 0,
             regex: None,
+            priority: None,
         })
     }
 
@@ -132,6 +198,16 @@ impl JournalStat {
         self
     }
 
+    /// Register interest in a particular log level.
+    pub fn set_filter_priority(&mut self, priority: &Option<String>) -> &mut Self {
+        if let Some(p) = priority {
+            let pstr = p.as_str();
+            self.priority = Some(pstr.into());
+        }
+
+        self
+    }
+
     /// Set the number of top talkers to watch for.
     pub fn n_frequent(&mut self, n_freq: usize) -> &mut Self {
         self.top_talkers = Vec::with_capacity(n_freq);
@@ -166,6 +242,13 @@ impl JournalStat {
                     }
                 }
 
+                let priority: Priority = priority.as_str().into();
+                if let Some(p) = &self.priority {
+                    if *p < priority {
+                        continue;
+                    }
+                }
+
                 self.total_msgs += 1;
 
                 let key = Message {
@@ -217,22 +300,6 @@ impl JournalStat {
         self
     }
 
-    /// Turn a number string priority into a syslog priority name.
-    fn pretty_priorty(&self, prio: &str) -> String {
-        match prio {
-            "0" => "emergency",
-            "1" => "alert",
-            "2" => "critical",
-            "3" => "error",
-            "4" => "warn",
-            "5" => "notice",
-            "6" => "info",
-            "7" => "debug",
-            _ => "unknown",
-        }
-        .to_string()
-    }
-
     /// Generate a report.
     pub fn report(&self) {
         println!("Journal statistics for {}", self.input.display());
@@ -267,7 +334,7 @@ impl JournalStat {
                     Rank: i + 1,
                     Frequency: *count,
                     Process: &msg.process,
-                    Priority: self.pretty_priorty(&msg.priority),
+                    Priority: &msg.priority.as_str(),
                     Message: &msg.msg,
                 });
             }
@@ -301,6 +368,7 @@ fn main() {
         .n_frequent(opt.top_talkers.unwrap_or(0))
         .n_largest(opt.large_messages.unwrap_or(0))
         .set_filter_unit(&opt.unit)
+        .set_filter_priority(&opt.priority)
         .set_regex(
             &opt.pattern
                 .map_or(None, |r| Some(Regex::new(&r).expect("invalid regex"))),
@@ -308,3 +376,16 @@ fn main() {
         .parse()
         .report();
 }
+
+#[cfg(test)]
+mod test {
+    use crate::Priority;
+
+    #[test]
+    fn test_priorities() {
+        let p0 = Priority::INFO;
+        let p1 = Priority::DEBUG;
+
+        assert!(p0 < p1);
+    }
+}