From 9c3e6fb886c239b81af8ee48080d9ea256178a5b Mon Sep 17 00:00:00 2001
From: Timothy <yexiaozhou2003@gmail.com>
Date: Fri, 1 Mar 2024 01:36:36 +0800
Subject: [PATCH] integrate with the server

---
 src/main.rs                   |  8 +++++++-
 src/sequence/mod.rs           |  2 ++
 src/sequence/port_sequence.rs | 13 +++++++++++--
 src/server/server.rs          | 19 +++++++++++--------
 4 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 16cc4e0..dc336f7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,4 @@
+use sequence::PortSequenceDetector;
 use server::Server;
 
 mod config;
@@ -5,7 +6,12 @@ mod sequence;
 mod server;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let server = Server::new("enp3s0".to_string());
+    // Load the configuration
+    let config = config::load_config("config.yaml")?;
+    // Create the sequence detector
+    let detector = PortSequenceDetector::new(config);
+
+    let mut server = Server::new("enp3s0".to_string(), Box::new(detector));
     server.start();
 
     Ok(())
diff --git a/src/sequence/mod.rs b/src/sequence/mod.rs
index 3455caf..2e97522 100644
--- a/src/sequence/mod.rs
+++ b/src/sequence/mod.rs
@@ -1,3 +1,5 @@
+pub use port_sequence::PortSequenceDetector;
+
 mod port_sequence;
 
 pub trait SequenceDetector {
diff --git a/src/sequence/port_sequence.rs b/src/sequence/port_sequence.rs
index 07d4562..6cdce06 100644
--- a/src/sequence/port_sequence.rs
+++ b/src/sequence/port_sequence.rs
@@ -41,8 +41,17 @@ impl SequenceDetector for PortSequenceDetector {
             return;
         }
 
-        let client_sequence = self.client_sequences.entry(client_ip).or_insert(Vec::new());
+        println!(
+            "SYN packet detected from: {} to target port: {}",
+            client_ip, sequence
+        );
+
+        let client_sequence = self
+            .client_sequences
+            .entry(client_ip.clone())
+            .or_insert(Vec::new());
         client_sequence.push(sequence);
+        self.match_sequence(&client_ip);
     }
 
     fn match_sequence(&mut self, client_ip: &str) -> bool {
@@ -51,7 +60,7 @@ impl SequenceDetector for PortSequenceDetector {
         if let Some(sequence) = client_sequence {
             for rule in &self.sequence_rules {
                 if sequence.ends_with(rule) {
-                    println!("Matched sequence: {:?}", rule);
+                    println!("Matched knock sequence: {:?} from: {}", rule, client_ip);
                     // clear the sequence
                     sequence.clear();
                     return true;
diff --git a/src/server/server.rs b/src/server/server.rs
index 9a8d519..e3bb83c 100644
--- a/src/server/server.rs
+++ b/src/server/server.rs
@@ -7,18 +7,22 @@ use pnet::packet::ip::IpNextHeaderProtocols;
 use pnet::packet::tcp::TcpPacket;
 use pnet::packet::Packet;
 
+use crate::sequence::SequenceDetector;
+
 pub struct Server {
     interface_name: String,
+    detector: Box<dyn SequenceDetector>,
 }
 
 impl Server {
-    pub fn new(interface: String) -> Server {
-        Server {
+    pub fn new(interface: String, detector: Box<dyn SequenceDetector>) -> Box<Server> {
+        Box::new(Server {
             interface_name: interface,
-        }
+            detector,
+        })
     }
 
-    pub fn start(&self) {
+    pub fn start(&mut self) {
         let interface = datalink::interfaces()
             .into_iter()
             .find(|iface: &NetworkInterface| iface.name == self.interface_name)
@@ -29,7 +33,7 @@ impl Server {
             Ok(Ethernet(tx, rx)) => (tx, rx),
             Ok(_) => panic!("Unhandled channel type"),
             Err(e) => panic!(
-                "An error occurred when creating the datalink channel: {}",
+                "An error occurred when creating the data link channel: {}",
                 e
             ),
         };
@@ -50,10 +54,9 @@ impl Server {
                                             && tcp.get_flags() & pnet::packet::tcp::TcpFlags::ACK
                                                 == 0
                                         {
-                                            println!(
-                                                "SYN packet detected from: {} to target port: {:?}",
+                                            self.detector.add_sequence(
                                                 header.get_source().to_string(),
-                                                tcp.get_destination()
+                                                tcp.get_destination() as i32,
                                             );
                                         }
                                     }