check client timeout

TimothyYe · Mar 1, 2024 · 12dd8c6 · 12dd8c6
1 parent b7622bf
commit 12dd8c6
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 7 deletions.
diff --git a/src/config/config.rs b/src/config/config.rs
@@ -12,6 +12,6 @@ pub struct Rule {
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Config {
     pub interface: String,
-    pub timeout: i32,
+    pub timeout: u64,
     pub rules: Vec<Rule>,
 }
diff --git a/src/config/mod.rs b/src/config/mod.rs
@@ -2,9 +2,7 @@ use std::fs::File;
 use std::io::Read;
 
 pub use config::Config;
-pub use config::Rule;
-
-mod config;
+pub mod config;
 
 pub fn load_config(path: &str) -> Result<Config, Box<dyn std::error::Error>> {
     let mut file = File::open(path)?;

diff --git a/src/sequence/port_sequence.rs b/src/sequence/port_sequence.rs
@@ -1,17 +1,21 @@
 use std::collections::{HashMap, HashSet};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use crate::config::Config;
+
 use crate::sequence::SequenceDetector;
 
 #[derive(Debug)]
 pub struct PortSequenceDetector {
-    timeout: i32,
+    timeout: u64,
     sequence_set: HashSet<i32>,
     sequence_rules: Vec<Vec<i32>>,
     client_sequences: HashMap<String, Vec<i32>>,
+    client_timeout: HashMap<String, u64>,
 }
 
 impl PortSequenceDetector {
+    #[must_use]
     pub fn new(config: Config) -> PortSequenceDetector {
         let mut sequence_rules = Vec::new();
         for rule in config.rules.clone() {
@@ -30,6 +34,7 @@ impl PortSequenceDetector {
             sequence_set,
             sequence_rules,
             client_sequences: HashMap::new(),
+            client_timeout: HashMap::new(),
         }
     }
 }
@@ -51,6 +56,15 @@ impl SequenceDetector for PortSequenceDetector {
             .entry(client_ip.clone())
             .or_insert(Vec::new());
         client_sequence.push(sequence);
+
+        // get the current time stamp
+        self.client_timeout.entry(client_ip.clone()).or_insert(
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_secs(),
+        );
+
         self.match_sequence(&client_ip);
     }
 
@@ -66,6 +80,21 @@ impl SequenceDetector for PortSequenceDetector {
                     return true;
                 }
             }
+
+            // check if the sequence has expired
+            let timeout_entry = self.client_timeout.get(client_ip);
+            if let Some(timeout) = timeout_entry {
+                let current_time = SystemTime::now()
+                    .duration_since(UNIX_EPOCH)
+                    .unwrap()
+                    .as_secs();
+
+                if current_time - timeout > self.timeout {
+                    println!("Sequence timeout for: {}", client_ip);
+                    sequence.clear();
+                    self.client_timeout.remove(client_ip);
+                }
+            }
         }
 
         false
@@ -81,12 +110,12 @@ mod tests {
             interface: "enp3s0".to_string(),
             timeout: 5,
             rules: vec![
-                crate::config::Rule {
+                crate::config::config::Rule {
                     name: "enable ssh".to_string(),
                     sequence: vec![1, 2, 3],
                     command: "ls -lh".to_string(),
                 },
-                crate::config::Rule {
+                crate::config::config::Rule {
                     name: "disable ssh".to_string(),
                     sequence: vec![3, 5, 6],
                     command: "du -sh *".to_string(),