Move logic into traits

src/main.rs

@@ -31,8 +31,8 @@ fn main() {
         Some(pf) => Capture::from_file(pf).expect("Invalid pcap file provided"),
         None => {
             warn!("Using example.pcap as no pcap was provided");
-            Capture::from_file(PathBuf::from("./pcaps/ssh_test.pcap"))
-                .expect("./pcaps/ssh_test.pcap does not exist")
+            Capture::from_file(PathBuf::from("./pcaps/dns.cap"))
+                .expect("./pcaps/dns.cap does not exist")
         }
     };
     info!("Pcap loaded");
@@ -57,6 +57,9 @@ fn main() {
                             stats.known_packets += 1;
                             trace!("Known packet type found: {:?}", x);
                             let extracted = extract_info(x, sliced.payload.to_vec());
+                            if extracted.is_some() {
+                                stats.analyzed += 1;
+                            }
                             if opt.print_analysis {
                                 info!("{:?}", extracted);
                             }

src/protocols/dns.rs

@@ -1,4 +1,4 @@
-use crate::util::*;
+use crate::{util::*, ExtractedInfo, KnownProtocol, ProtocolType};
 use nom::{
     bytes::complete::{tag, take},
     combinator::map,
@@ -15,6 +15,22 @@ pub enum DNSType {
     Response,
 }
 
+impl KnownProtocol for DNSType {
+    fn extract_info(&self, payload: Vec<u8>) -> ExtractedInfo {
+        ExtractedInfo::DNSQuery(analyse_dns_query(payload))
+    }
+
+    fn classify_proto(payload: Vec<u8>) -> Result<ProtocolType, ()> {
+        if is_dns_query(&payload) {
+            return Ok(ProtocolType::DNS(DNSType::Query));
+        }
+        if is_dns_response(&payload) {
+            return Ok(ProtocolType::DNS(DNSType::Response));
+        }
+        Err(())
+    }
+}
+
 #[derive(Debug, PartialEq, Eq)]
 pub struct DNSValue {
     pub txid: u16,

src/protocols/mod.rs

@@ -1,6 +1,12 @@
 use dns::{DNSType, DNSValue};
 
 mod dns;
+mod ssh;
+
+pub trait KnownProtocol {
+    fn classify_proto(payload: Vec<u8>) -> Result<ProtocolType, ()>;
+    fn extract_info(&self, payload: Vec<u8>) -> ExtractedInfo;
+}
 
 #[derive(Debug)]
 pub enum ProtocolType {
@@ -24,9 +30,8 @@ pub fn extract_info(ptype: ProtocolType, payload: Vec<u8>) -> Option<ExtractedIn
 }
 
 pub fn match_protocol(payload: Vec<u8>) -> Result<ProtocolType, ()> {
-    match dns::is_dns(payload) {
-        Ok(x) => return Ok(ProtocolType::DNS(x)),
-        Err(_) => {}
+    if let Ok(x) = dns::is_dns(payload) {
+        return Ok(ProtocolType::DNS(x));
     };
-    return Err(());
+    Err(())
 }

src/protocols/ssh.rs

@@ -1 +1,4 @@
-
+pub fn is_ssh(payload: Vec<u8>) -> bool {
+    // Check for ASCII "SSH"
+    payload[0] == 0x53 && payload[1] == 0x53 && payload[2] == 0x48
+}

src/util.rs

@@ -11,6 +11,7 @@ pub struct Stats {
     pub unknown_packets: usize,
     pub errored_packets: usize,
     pub empty_payload: usize,
+    pub analyzed: usize,
 }
 
 impl Stats {
@@ -21,6 +22,7 @@ impl Stats {
             unknown_packets: 0,
             errored_packets: 0,
             empty_payload: 0,
+            analyzed: 0,
         }
     }
 
@@ -37,9 +39,10 @@ impl Display for Stats {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
             f,
-            "Total: {}\nKnown: {}\nUnknown: {}\nErrored: {}\nEmpty: {}",
+            "Total: {}\nKnown/Analyzed: {}({})\nUnknown: {}\nErrored: {}\nEmpty: {}",
             self.total_packets,
             self.known_packets,
+            self.analyzed,
             self.unknown_packets,
             self.errored_packets,
             self.empty_payload