use std::io::{self, Write};
use std::net::{TcpStream, UdpSocket};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
#[derive(Clone, Debug)]
pub enum Protocol {
TCP,
UDP,
}
pub struct Pentest {
ip: String,
port: u16,
counter: usize,
protocol: Protocol,
payload: Vec<u8>,
concurrency: usize,
}
impl Pentest {
pub fn new(
ip: &str,
port: u16,
counter: usize,
protocol: Protocol,
payload: Vec<u8>,
concurrency: usize,
) -> Pentest {
Pentest {
ip: ip.to_string(),
port,
counter,
protocol,
payload,
concurrency,
}
}
pub fn run(&self) -> io::Result<()> {
let tasks = self.counter / self.concurrency;
let tasks_remainder = self.counter % self.concurrency;
let ip = Arc::new(self.ip.clone());
let payload = Arc::new(self.payload.clone());
let mut threads = vec![];
println!("Starting test with {} threads", self.concurrency);
for _ in 0..self.concurrency {
let ip = Arc::clone(&ip);
let payload = Arc::clone(&payload);
let pentest = Pentest {
ip: ip.to_string(),
port: self.port,
counter: tasks,
protocol: self.protocol.clone(),
payload: payload.to_vec(),
concurrency: 1,
};
threads.push(thread::spawn(move || {
pentest.run_inner().unwrap();
}));
}
if tasks_remainder > 0 {
let pentest = Pentest {
ip: self.ip.clone(),
port: self.port,
counter: tasks_remainder,
protocol: self.protocol.clone(),
payload: self.payload.clone(),
concurrency: 1,
};
pentest.run_inner()?;
}
for (i, thread) in threads.into_iter().enumerate() {
println!("Joining thread {}", i + 1);
match thread.join() {
Ok(_) => println!("Thread {} joined", i + 1),
Err(err) => eprintln!("Error in thread {}: {:?}", i + 1, err),
}
}
Ok(())
}
fn run_inner(&self) -> io::Result<()> {
match self.protocol {
Protocol::TCP => self.run_tcp(),
Protocol::UDP => self.run_udp(),
}
}
fn run_tcp(&self) -> io::Result<()> {
let addr = format!("{}:{}", self.ip, self.port);
for i in 0..self.counter {
println!("Sending TCP packet {} to {}:{}", i + 1, self.ip, self.port);
match TcpStream::connect(&addr) {
Ok(mut stream) => {
stream.set_write_timeout(Some(Duration::from_secs(1)))?;
stream.write_all(&self.payload)?;
}
Err(err) => eprintln!("Failed to send TCP packet {}: {}", i + 1, err),
}
}
println!("Finished sending TCP packets");
Ok(())
}
fn run_udp(&self) -> io::Result<()> {
let socket = UdpSocket::bind("0.0.0.0:0")?;
socket.set_write_timeout(Some(Duration::from_secs(1)))?;
let addr = format!("{}:{}", self.ip, self.port);
for i in 0..self.counter {
println!("Sending UDP packet {} to {}:{}", i + 1, self.ip, self.port);
if let Err(err) = socket.send_to(&self.payload, &addr) {
eprintln!("Failed to send UDP packet {}: {}", i + 1, err);
}
}
println!("Finished sending UDP packets");
Ok(())
}
}