added uniqueby

QpxDesign · May 14, 2024 · e5f373e · e5f373e
1 parent c996423
commit e5f373e
Show file tree

Hide file tree

Showing 16 changed files with 139,486 additions and 313,985 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,7 @@
 workspace = { members = [ "tests","tests/obfuscate_access_log_ips"] }
 [package]
 name = "ngxav"
-version = "0.5.5"
+version = "0.6.0"
 edition = "2021"
 license = "MIT"
 description = "Search through NGINX logs with advanced filters and support for displaying analytics about your selected log entries"
@@ -11,6 +11,7 @@ repository = "https://github.com/qpxdesign/ngxav-rs/"
 [dependencies]
 chrono = "0.4.33"
 clap = { version = "4.4.18", features = ["derive"] }
+fancy-regex = "0.13.0"
 flate2 = "1.0.28"
 isbot = "0.1.3"
 lazy_static = "1.4.0"

diff --git a/src/main.rs b/src/main.rs
@@ -81,19 +81,19 @@ fn main() {
             utils::keep_line::keep_line(p, false) == true
         })
         .collect();
-    if !args.unique.is_none() && args.unique == Some(true) {
-        kel = utils::unique_ips_only::unique_ips_only(kel);
+    if (!args.unique.is_none() && args.unique == Some(true)) {
+        kel = utils::unique_ips_only::unique_ips_only(kel, args.unique_by.clone());
     }
     if !args.analytics.is_none() && args.analytics == Some(true) {
         utils::generate_analytics::generate_analytical_output(kel);
     } else if !args.session_analytics.is_none() && args.session_analytics == Some(true) {
-        utils::session_analytics::session_analytics(kel);
+        utils::session_analytics::session_analytics(kel, args.unique_by.clone());
     } else if !args.large.is_none() {
         utils::sort_by_body_size::sort_by_body_size(kel, args.large.unwrap());
     } else if !args.ip_ses.is_none() {
-        utils::sessions_from_ip::sessions_from_ip(kel, args.ip_ses.unwrap());
+        utils::sessions_from_ip::sessions_from_ip(kel, args.ip_ses, args.unique_by);
     } else if !args.session_unqiue.is_none() && args.session_unqiue == Some(true) {
-        utils::session_unique::session_unique(kel);
+        utils::session_unique::session_unique(kel, args.unique_by);
     } else {
         kel.par_sort_by_key(|a| parse_nginx_time_format(&a.time).timestamp());
         for line in kel {

diff --git a/src/structs/Args.rs b/src/structs/Args.rs
@@ -71,4 +71,7 @@ pub struct ArgParser {
 
     #[arg(short = 'c', long = "cm", action=ArgAction::SetTrue)]
     pub conserve_memory: Option<bool>,
+
+    #[arg(short = 'y', long = "ub")]
+    pub unique_by: Option<String>,
 }
diff --git a/src/utils/session_analytics.rs b/src/utils/session_analytics.rs
@@ -14,8 +14,11 @@ struct SessionAnalysisStats {
     request_count_sum: i64,
     request_length_sum: i64,
 }
-pub fn session_analytics(log_selection: Vec<crate::structs::LineParseResult::LineParseResult>) {
-    let mut sessions = sessionize(log_selection);
+pub fn session_analytics(
+    log_selection: Vec<crate::structs::LineParseResult::LineParseResult>,
+    unique_by: Option<String>,
+) {
+    let mut sessions = sessionize(log_selection, unique_by);
     let mut stats: SessionAnalysisStats = SessionAnalysisStats {
         total_count: 0,
         host_paths: HashMap::new(),

diff --git a/src/utils/session_unique.rs b/src/utils/session_unique.rs
@@ -3,8 +3,11 @@ use crate::utils::parse_nginx_time_format::parse_nginx_time_format;
 use crate::utils::sessionize::sessionize;
 use rayon::slice::ParallelSliceMut;
 
-pub fn session_unique(log_selection: Vec<crate::structs::LineParseResult::LineParseResult>) {
-    let sessions = sessionize(log_selection);
+pub fn session_unique(
+    log_selection: Vec<crate::structs::LineParseResult::LineParseResult>,
+    unique_by: Option<String>,
+) {
+    let sessions = sessionize(log_selection, unique_by);
     let mut out: Vec<String> = Vec::new();
     for ip_session in sessions {
         for session in ip_session.sessions {

diff --git a/src/utils/sessionize.rs b/src/utils/sessionize.rs
@@ -1,6 +1,6 @@
 use crate::utils::parse_nginx_time_format::parse_nginx_time_format;
 use chrono::prelude::*;
-use std::borrow::Cow;
+use regex::Regex;
 use std::collections::HashMap;
 
 #[derive(Clone)]
@@ -10,41 +10,53 @@ pub struct SessionOccurrences {
     pub times: Vec<DateTime<Utc>>,
     pub sessions: Vec<Vec<String>>,
 }
+
 pub fn sessionize(
     lines: Vec<crate::structs::LineParseResult::LineParseResult>,
+    unique_by: Option<String>,
 ) -> Vec<SessionOccurrences> {
     let session_cutoff_min = 10;
-    let mut occurrences: HashMap<&str, SessionOccurrences> = HashMap::new();
-
+    let mut occurrences: HashMap<String, SessionOccurrences> = HashMap::new();
+    let mut r = Regex::new("").unwrap();
+    if unique_by.is_some() {
+        let u = &unique_by.clone().unwrap();
+        r = Regex::new(&u).unwrap();
+    }
     for parsed_line in lines {
-        if parsed_line.ip_address != "-" {
+        let mut key: String = "".to_string();
+        if unique_by.is_some() {
+            let m = r.find(parsed_line.full_text);
+            if m.is_some() {
+                key = m.unwrap().as_str().to_string();
+            } else {
+                continue;
+            }
+        } else {
+            key = parsed_line.ip_address.to_string()
+        }
+        if &key != "-" {
             let time: DateTime<Utc> = parse_nginx_time_format(&parsed_line.time);
-            if !occurrences.contains_key(&parsed_line.ip_address) {
-                let cl = parsed_line.ip_address.to_owned();
+            if !occurrences.contains_key(&key) {
                 let mut l = Vec::new();
                 l.push(parsed_line.full_text.to_string());
                 let mut t = Vec::new();
                 t.push(time);
                 occurrences.insert(
-                    parsed_line.ip_address,
+                    key.clone(),
                     SessionOccurrences {
-                        ip_address: cl,
+                        ip_address: key.clone(),
                         lines: l,
                         times: t,
                         sessions: Vec::new(),
                     },
                 );
             } else {
                 occurrences
-                    .get_mut(&parsed_line.ip_address)
+                    .get_mut(key.as_str())
                     .unwrap()
                     .lines
                     .push(parsed_line.full_text.to_string());
-                occurrences
-                    .get_mut(&parsed_line.ip_address)
-                    .unwrap()
-                    .times
-                    .push(time);
+                occurrences.get_mut(key.as_str()).unwrap().times.push(time);
             }
         }
     }

diff --git a/src/utils/sessions_from_ip.rs b/src/utils/sessions_from_ip.rs
@@ -1,15 +1,31 @@
 use crate::utils::parse_line::parse_line;
 use crate::utils::sessionize::sessionize;
+use regex::Regex;
+
 pub fn sessions_from_ip(
     log_selection: Vec<crate::structs::LineParseResult::LineParseResult>,
-    ip: String,
+    ip: Option<String>,
+    unique_by: Option<String>,
 ) {
-    let sessions = sessionize(log_selection);
+    let sessions = sessionize(log_selection, unique_by.clone());
     let mut host_paths: Vec<Vec<String>> = Vec::new();
     let mut session_start_times: Vec<String> = Vec::new();
     let mut session_end_times: Vec<String> = Vec::new();
     for session_entry in sessions {
-        if session_entry.ip_address == ip {
+        let mut key: String = "".to_string();
+        if ip.clone().is_some() {
+            key = ip.clone().unwrap();
+        } else if unique_by.clone().is_some() {
+            let u = unique_by.clone().unwrap();
+            let r = Regex::new(&regex::escape(&u)).unwrap();
+            let m = r.find(&session_entry.lines[0]);
+            if m.is_some() {
+                key = m.unwrap().clone().as_str().to_string();
+            } else {
+                continue;
+            }
+        }
+        if session_entry.ip_address == key {
             for session in session_entry.sessions {
                 let mut host_path: Vec<String> = Vec::new();
                 if session.len() != 0 && session.len() != 1 {

diff --git a/src/utils/unique_ips_only.rs b/src/utils/unique_ips_only.rs
@@ -1,14 +1,30 @@
 #[path = "../structs/mod.rs"]
 mod structs;
 use crate::structs::LineParseResult::LineParseResult;
+use lazy_static::lazy_static;
+use regex::Regex;
 use std::collections::HashMap;
 
-pub fn unique_ips_only(parsed_lines: Vec<LineParseResult>) -> Vec<LineParseResult> {
+pub fn unique_ips_only(
+    parsed_lines: Vec<LineParseResult>,
+    unique_by: Option<String>,
+) -> Vec<LineParseResult> {
     let mut occurrences: HashMap<&str, LineParseResult> = HashMap::new();
-
+    let mut r = Regex::new("").unwrap();
+    if unique_by.is_some() {
+        r = Regex::new(&unique_by.clone().unwrap()).unwrap();
+    }
     for line in parsed_lines {
-        if occurrences.contains_key(&line.ip_address) == false {
-            occurrences.insert(line.ip_address, line);
+        let mut u: String = "".to_string();
+        if unique_by.is_some() {
+            let m = r.find(line.full_text);
+            if m.is_some() {
+                occurrences.insert(m.unwrap().as_str(), line);
+            }
+        } else {
+            if occurrences.contains_key(&line.ip_address) == false {
+                occurrences.insert(line.ip_address, line);
+            }
         }
     }
     return occurrences.into_iter().map(|(_, v)| v).collect();