upd: add protocol arg

Author: teo-lohrer-su

Cargo.toml

@@ -11,7 +11,6 @@ caracat = { git = "https://github.com/maxmouchet/caracat", branch = "main" }
 chrono = "0.4.38"
 clap = { version = "4.5.16", features = ["derive"] }
 env_logger = "0.11.5"
-# caracat = "1.0.0"
 ip_network = "0.4.1"
 itertools = "0.12.0"
 log = "0.4.21"

src/algorithms/diamond_miner.rs

@@ -4,8 +4,8 @@ mod types;
 
 use caracat::models::{Probe, Reply, L4};
 use itertools::Itertools;
-use log::{debug, warn};
-// use log::debug;
+use log::debug;
+// use log::{debug, warn};
 pub use sequential_mapper::*;
 
 use std::collections::{HashMap, HashSet};
@@ -125,6 +125,8 @@ impl DiamondMiner {
             .map(|&node| (node, node_replies(&self.replies(), node, ttl)))
             .collect();
 
+        debug!("TTL {}: Link_dist: {:?}", ttl, link_dist);
+
         let total: usize = link_dist.values().sum();
 
         let link_dist: HashMap<IpAddr, f64> = if total > 0 {

src/main.rs

@@ -1,5 +1,4 @@
 use std::collections::{HashMap, HashSet};
-use std::io::Stdout;
 use std::net::{IpAddr, Ipv4Addr};
 use std::time::Duration;
 use std::{fmt, vec};
@@ -45,6 +44,21 @@ impl fmt::Display for OutputFormat {
     }
 }
 
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)]
+enum ProtocolArg {
+    ICMP,
+    UDP,
+}
+
+impl From<ProtocolArg> for caracat::models::L4 {
+    fn from(protocol: ProtocolArg) -> Self {
+        match protocol {
+            ProtocolArg::UDP => caracat::models::L4::UDP,
+            ProtocolArg::ICMP => caracat::models::L4::ICMP,
+        }
+    }
+}
+
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {
@@ -92,6 +106,10 @@ struct Args {
     #[arg(long, default_value_t = 100)]
     probing_rate: u64,
 
+    /// Protocol to use (ICMP or UDP)
+    #[arg(short, long, value_enum, default_value_t = ProtocolArg::ICMP)]
+    protocol: ProtocolArg,
+
     /// Network interface to use
     #[arg(short, long)]
     interface: Option<String>,
@@ -112,7 +130,7 @@ fn main() -> Result<()> {
     let max_ttl = args.max_ttl;
     let src_port = args.src_port;
     let dst_port = args.dst_port;
-    let protocol = caracat::models::L4::ICMP;
+    let protocol = args.protocol.into();
     let confidence = args.confidence;
     let max_round = args.max_round;
     let estimate_successsors = args.estimate_successors;
@@ -264,7 +282,7 @@ fn main() -> Result<()> {
         OutputFormat::Atlas => {
             debug!("--- ATLAS output ---");
             let stdout = std::io::stdout();
-            let mut atlas_writer: AtlasWriter<Stdout> = AtlasWriter::new(stdout);
+            let mut atlas_writer = AtlasWriter::new(stdout);
             atlas_writer.write_traceroute(&traceroute)?;
         }
         OutputFormat::Iris => {