feature: added allowed_versions. (#104)

* feature: added allowed_versions. * fix: Added snap tests --------- Co-authored-by: mikemiles-dev <michaelmileusnich@Michaels-MacBook-Air-2.local>
mikemiles-dev · Jan 27, 2025 · 02285a4 · 02285a4
1 parent 90a110a
commit 02285a4
Show file tree

Hide file tree

Showing 9 changed files with 148 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -48,6 +48,19 @@ let parsed = NetflowParser::default().parse_bytes(&v5_packet);
 let v5_parsed: Vec<NetflowPacket> = parsed.into_iter().filter(|p| p.is_v5()).collect();
 ```
 
+## Parsing out uneeded versions
+If you only care about a specific version or versions you can specfic `allowed_version`:
+```rust
+use netflow_parser::{NetflowParser, NetflowPacket};
+
+let v5_packet = [0, 5, 0, 1, 3, 0, 4, 0, 5, 0, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7,];
+let mut parser = NetflowParser::default();
+parser.allowed_versions = [7, 9].into();
+let parsed = NetflowParser::default().parse_bytes(&v5_packet);
+```
+
+This code will return an empty Vec as version 5 is not allowed.
+
 ## Netflow Common
 
 We have included a `NetflowCommon` and `NetflowCommonFlowSet` structure.

diff --git a/RELEASES.md b/RELEASES.md
@@ -1,5 +1,6 @@
 # 0.5.1
 * Reworked NetflowParseError.  Added a Partial Type.
+* Added ability to parse only `allowed_versions`.
 * V9, IPFix, Datanumber Code cleanup.
 * Added benchmarking
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -48,6 +48,19 @@
 //! let v5_parsed: Vec<NetflowPacket> = parsed.into_iter().filter(|p| p.is_v5()).collect();
 //! ```
 //!
+//! ## Parsing out uneeded versions
+//! If you only care about a specific version or versions you can specfic `allowed_version`:
+//! ```rust
+//! use netflow_parser::{NetflowParser, NetflowPacket};
+//!
+//! let v5_packet = [0, 5, 0, 1, 3, 0, 4, 0, 5, 0, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7,];
+//! let mut parser = NetflowParser::default();
+//! parser.allowed_versions = [7, 9].into();
+//! let parsed = NetflowParser::default().parse_bytes(&v5_packet);
+//! ```
+//!
+// !This code will return an empty Vec as version 5 is not allowed.
+//!
 //! ## Netflow Common
 //!
 //! We have included a `NetflowCommon` and `NetflowCommonFlowSet` structure.
@@ -184,6 +197,8 @@ use crate::variable_versions::v9;
 use nom_derive::{Nom, Parse};
 use serde::Serialize;
 
+use std::collections::HashSet;
+
 /// Enum of supported Netflow Versions
 #[derive(Debug, Clone, Serialize)]
 pub enum NetflowPacket {
@@ -226,10 +241,11 @@ struct GenericNetflowHeader {
     version: u16,
 }
 
-#[derive(Default, Debug)]
+#[derive(Debug)]
 pub struct NetflowParser {
     pub v9_parser: V9Parser,
     pub ipfix_parser: IPFixParser,
+    pub allowed_versions: HashSet<u16>,
 }
 
 #[derive(Debug, Clone)]
@@ -258,6 +274,7 @@ pub struct NetflowPacketError {
 pub enum NetflowParseError {
     Incomplete(String),
     Partial(PartialParse),
+    UnallowedVersion(u16),
     UnknownVersion(Vec<u8>),
 }
 
@@ -268,6 +285,16 @@ pub struct PartialParse {
     pub error: String,
 }
 
+impl Default for NetflowParser {
+    fn default() -> Self {
+        Self {
+            v9_parser: V9Parser::default(),
+            ipfix_parser: IPFixParser::default(),
+            allowed_versions: [5, 7, 9, 10].iter().cloned().collect(),
+        }
+    }
+}
+
 impl NetflowParser {
     /// Takes a Netflow packet slice and returns a vector of Parsed Netflows.
     /// If we reach some parse error we return what items be have.
@@ -302,10 +329,29 @@ impl NetflowParser {
                 }
                 results
             }
-            Err(e) => vec![NetflowPacket::Error(NetflowPacketError {
-                error: e,
-                remaining: packet.to_vec(),
-            })],
+            Err(e) => match e {
+                NetflowParseError::Incomplete(_) => {
+                    vec![NetflowPacket::Error(NetflowPacketError {
+                        error: e,
+                        remaining: packet.to_vec(),
+                    })]
+                }
+                NetflowParseError::Partial(partial) => {
+                    vec![NetflowPacket::Error(NetflowPacketError {
+                        error: NetflowParseError::Partial(partial),
+                        remaining: packet.to_vec(),
+                    })]
+                }
+                NetflowParseError::UnknownVersion(_) => {
+                    vec![NetflowPacket::Error(NetflowPacketError {
+                        error: e,
+                        remaining: packet.to_vec(),
+                    })]
+                }
+                NetflowParseError::UnallowedVersion(_) => {
+                    vec![]
+                }
+            },
         }
     }
 
@@ -333,6 +379,10 @@ impl NetflowParser {
             .map(|(remaining, header)| (remaining, header.version))
             .map_err(|e| NetflowParseError::Incomplete(e.to_string()))?;
 
+        if !self.allowed_versions.contains(&version) {
+            return Err(NetflowParseError::UnallowedVersion(version));
+        }
+
         match version {
             5 => v5::parse_netflow_v5(packet),
             7 => v7::parse_netflow_v7(packet),

diff --git a/src/snapshots/netflow_parser__tests__base_tests__it_creates_error.snap b/src/snapshots/netflow_parser__tests__base_tests__it_creates_error.snap
@@ -4,10 +4,14 @@ expression: "NetflowParser::default().parse_bytes(&packet)"
 ---
 - Error:
     error:
-      UnknownVersion:
-        - 14
+      Partial:
+        version: 9
+        remaining:
+          - 10
+          - 11
+        error: Parsing requires 4 bytes/chars
     remaining:
-      - 12
-      - 13
-      - 14
-
+      - 0
+      - 9
+      - 10
+      - 11
diff --git a/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_ipfix.snap b/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_ipfix.snap
@@ -0,0 +1,5 @@
+---
+source: src/tests.rs
+expression: parser.parse_bytes(&packet)
+---
+[]
diff --git a/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_v5.snap b/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_v5.snap
@@ -0,0 +1,5 @@
+---
+source: src/tests.rs
+expression: parser.parse_bytes(&packet)
+---
+[]
diff --git a/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_v7.snap b/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_v7.snap
@@ -0,0 +1,5 @@
+---
+source: src/tests.rs
+expression: parser.parse_bytes(&packet)
+---
+[]
diff --git a/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_v9.snap b/src/snapshots/netflow_parser__tests__base_tests__it_doesnt_allow_v9.snap
@@ -0,0 +1,5 @@
+---
+source: src/tests.rs
+expression: parser.parse_bytes(&packet)
+---
+[]
diff --git a/src/tests.rs b/src/tests.rs
@@ -11,6 +11,7 @@ mod base_tests {
 
     use hex;
     use insta::assert_yaml_snapshot;
+    use std::collections::HashSet;
 
     #[test]
     fn it_parses_unix_timestamp_correctly() {
@@ -48,6 +49,18 @@ mod base_tests {
         assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
     }
 
+    #[test]
+    fn it_doesnt_allow_v5() {
+        let packet = [
+            0, 5, 0, 1, 3, 0, 4, 0, 5, 0, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3,
+            4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1,
+            2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7,
+        ];
+        let mut parser = NetflowParser::default();
+        parser.allowed_versions = HashSet::default();
+        assert_yaml_snapshot!(parser.parse_bytes(&packet));
+    }
+
     #[test]
     fn it_parses_v5_incomplete() {
         let packet = [0, 5, 0, 0, 1, 1, 1, 1];
@@ -73,7 +86,7 @@ mod base_tests {
 
     #[test]
     fn it_creates_error() {
-        let packet = [12, 13, 14];
+        let packet = [0, 9, 10, 11];
         assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
     }
 
@@ -87,6 +100,18 @@ mod base_tests {
         assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
     }
 
+    #[test]
+    fn it_doesnt_allow_v7() {
+        let packet = [
+            0, 7, 0, 1, 3, 0, 4, 0, 5, 0, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3,
+            4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1,
+            2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1,
+        ];
+        let mut parser = NetflowParser::default();
+        parser.allowed_versions = HashSet::default();
+        assert_yaml_snapshot!(parser.parse_bytes(&packet));
+    }
+
     #[test]
     fn it_parses_v7_and_re_exports() {
         let packet = [
@@ -113,6 +138,17 @@ mod base_tests {
         assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
     }
 
+    #[test]
+    fn it_doesnt_allow_v9() {
+        let packet = [
+            0, 9, 0, 2, 0, 0, 9, 9, 0, 1, 2, 3, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 16, 1, 2, 0,
+            2, 0, 1, 0, 4, 0, 8, 0, 4, 1, 2, 0, 12, 9, 2, 3, 4, 9, 9, 9, 8,
+        ];
+        let mut parser = NetflowParser::default();
+        parser.allowed_versions = HashSet::default();
+        assert_yaml_snapshot!(parser.parse_bytes(&packet));
+    }
+
     #[test]
     fn it_parses_v9_ipv6flowlabel() {
         let templates_hex = r#"0009000400a21e176658cb4600000155000000080000004c0102001100080004000c0004000f000400070002000b0002000a0002000e000200fc000400fd000400020004000100040016000400150004000400010005000101000002003d0001000000540103001300080004000c0004000f000400070002000b000200060001000a0002000e000200fc000400fd000400020004000100040016000400150004000400010005000100d1000801000002003d00010000005401050013001b0010001c0010003e001000070002000b000200060001000a0002000e000200fc000400fd00040002000400010004001600040015000400040001000500010050000601000002003d00010000005801060014001b0010001c0010003e0010001f000300070002000b000200060001000a0002000e000200fc000400fd00040002000400010004001600040015000400040001000500010050000601000002003d0001"#;
@@ -243,6 +279,18 @@ mod base_tests {
         assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
     }
 
+    #[test]
+    fn it_doesnt_allow_ipfix() {
+        let packet = [
+            0, 10, 0, 64, 1, 2, 3, 4, 0, 0, 0, 0, 1, 2, 3, 4, 0, 2, 0, 20, 1, 0, 0, 3, 0, 8, 0,
+            4, 0, 12, 0, 4, 0, 2, 0, 4, 1, 0, 0, 28, 1, 2, 3, 4, 1, 2, 3, 3, 1, 2, 3, 2, 0, 2,
+            0, 2, 0, 1, 2, 3, 4, 5, 6, 7,
+        ];
+        let mut parser = NetflowParser::default();
+        parser.allowed_versions = HashSet::default();
+        assert_yaml_snapshot!(parser.parse_bytes(&packet));
+    }
+
     #[test]
     fn it_parses_ipfix_and_re_exports() {
         let packet = [