From 5726a604ae5ea78c3e929b98757e1f02cb1ee5e4 Mon Sep 17 00:00:00 2001
From: SkuldNorniern <skuldnorniern@gmail.com>
Date: Mon, 8 Jan 2024 22:50:22 +0900
Subject: [PATCH] fix: correct `IndexError`'s var into `usize`

---
 src/net/capture.rs     | 46 +++++++++++++++++++++++-------------------
 src/net/mod.rs         |  8 ++++----
 src/net/packet_pcap.rs | 14 +++++--------
 3 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/src/net/capture.rs b/src/net/capture.rs
index 0da63e1..abae835 100644
--- a/src/net/capture.rs
+++ b/src/net/capture.rs
@@ -1,24 +1,29 @@
 use pcap::{Active, Address, Capture, Device, Error as PcapError};
+
 use std::time::Instant;
+use std::fmt;
 
 #[derive(Debug)]
-pub enum CaptureError {
+pub enum DeviceError {
     Cap(PcapError),
-    Other(String),
     DeviceNotFound(String),
-    InitializationError(String),
-    InvalidDeviceIndex(String),
+    InitializationError(),
+    InvalidDeviceIndex(usize),
 }
 
-impl From<pcap::Error> for CaptureError {
+impl From<PcapError> for DeviceError {
     fn from(err: PcapError) -> Self {
-        CaptureError::Cap(err)
+        DeviceError::Cap(err)
     }
 }
-
-impl From<String> for CaptureError {
-    fn from(err: String) -> Self {
-        CaptureError::Other(err)
+impl fmt::Display for DeviceError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            DeviceError::Cap(err) => err.fmt(f),
+            DeviceError::DeviceNotFound(err) => write!(f, "Device not found: {}", err),
+            DeviceError::InitializationError() => write!(f, "Initialization error"),
+            DeviceError::InvalidDeviceIndex(err) => write!(f, "Invalid device index: {}", err),
+        }
     }
 }
 
@@ -30,7 +35,7 @@ pub struct CaptureDevice {
 }
 
 impl CaptureDevice {
-    pub fn new(device: Device) -> Result<CaptureDevice, CaptureError> {
+    pub fn new(device: Device) -> Result<CaptureDevice, DeviceError> {
         let capture = initialize_capture(device.clone())?;
 
         Ok(CaptureDevice {
@@ -42,11 +47,11 @@ impl CaptureDevice {
     }
 }
 
-pub fn list_devices() -> Result<Vec<Device>, CaptureError> {
-    Device::list().map_err(CaptureError::Cap)
+pub fn list_devices() -> Result<Vec<Device>, DeviceError> {
+    Device::list().map_err(DeviceError::Cap)
 }
 
-pub fn find_device(identifier: &str) -> Result<Device, CaptureError> {
+pub fn find_device(identifier: &str) -> Result<Device, DeviceError> {
     let start = Instant::now();
     println!("Requested Device: {}", identifier);
 
@@ -58,7 +63,7 @@ pub fn find_device(identifier: &str) -> Result<Device, CaptureError> {
             println!("Device {} captured in {:?}", device.name, duration);
             return Ok(device.clone());
         } else {
-            return Err(CaptureError::InvalidDeviceIndex(identifier.to_string()));
+            return Err(DeviceError::InvalidDeviceIndex(index));
         }
     }
 
@@ -70,16 +75,15 @@ pub fn find_device(identifier: &str) -> Result<Device, CaptureError> {
         }
     }
 
-    Err(CaptureError::DeviceNotFound(identifier.to_string()))
+    Err(DeviceError::DeviceNotFound(identifier.to_string()))
 }
 
-fn initialize_capture(device: Device) -> Result<Capture<Active>, CaptureError> {
-    Capture::from_device(device)
-        .map_err(CaptureError::Cap)?
+fn initialize_capture(device: Device) -> Result<Capture<Active>, DeviceError> {
+    Ok(Capture::from_device(device)
+        .map_err(DeviceError::Cap)?
         .promisc(true)
         .snaplen(1024)
         .timeout(60000)
         .immediate_mode(true)
-        .open()
-        .map_err(|_| CaptureError::InitializationError("Capture initialization error".into()))
+        .open()?)
 }
diff --git a/src/net/mod.rs b/src/net/mod.rs
index 34fe534..c15a9e0 100644
--- a/src/net/mod.rs
+++ b/src/net/mod.rs
@@ -2,7 +2,7 @@ pub mod errors;
 //mod fluereflow;
 mod capture;
 mod flows;
-mod interface;
+// mod interface;
 pub mod live_fluereflow;
 mod offline_fluereflows;
 pub mod online_fluereflow;
@@ -14,9 +14,9 @@ pub mod types;
 pub use capture::find_device;
 pub use capture::list_devices;
 pub use capture::CaptureDevice;
-pub use capture::CaptureError;
-pub use interface::list_interface_names;
-pub use interface::list_interfaces;
+pub use capture::DeviceError;
+// pub use interface::list_interface_names;
+// pub use interface::list_interfaces;
 pub use offline_fluereflows::fluereflow_fileparse;
 pub use packet_pcap::pcap_capture;
 //pub use types::FluereRecord;
diff --git a/src/net/packet_pcap.rs b/src/net/packet_pcap.rs
index 15f4c16..f2c9aa2 100644
--- a/src/net/packet_pcap.rs
+++ b/src/net/packet_pcap.rs
@@ -1,13 +1,13 @@
 extern crate chrono;
 
-use pcap::Capture;
 use std::fs;
 
 use tokio::time::sleep;
 
-use super::interface::get_interface;
 use crate::types::Args;
 use crate::utils::cur_time_file;
+use crate::net::CaptureDevice;
+use crate::net::find_device;
 
 use std::time::{Duration, Instant};
 
@@ -19,13 +19,9 @@ pub async fn pcap_capture(args: Args) {
     let sleep_windows = args.parameters.sleep_windows.unwrap();
     let verbose = args.verbose.unwrap();
 
-    let interface = get_interface(interface_name.as_str());
-    let mut cap = Capture::from_device(interface)
-        .unwrap()
-        .promisc(true)
-        //.buffer_size(1000000000)
-        .open()
-        .unwrap();
+    let interface = find_device(interface_name.as_str()).unwrap();
+    let cap_device = CaptureDevice::new(interface.clone()).unwrap();
+    let mut cap = cap_device.capture;
 
     let file_dir = "./output";
     let mut packet_count = 0;