pentest/
lib.rs

1use std::io::{self, Write};
2use std::net::{TcpStream, UdpSocket};
3use std::sync::Arc;
4use std::thread;
5use std::time::Duration;
6
7#[derive(Clone, Debug)]
8pub enum Protocol {
9    TCP,
10    UDP,
11}
12
13pub struct Pentest {
14    ip: String,
15    port: u16,
16    counter: usize,
17    protocol: Protocol,
18    payload: Vec<u8>,
19    concurrency: usize,
20}
21
22impl Pentest {
23    pub fn new(
24        ip: &str,
25        port: u16,
26        counter: usize,
27        protocol: Protocol,
28        payload: Vec<u8>,
29        concurrency: usize,
30    ) -> Pentest {
31        Pentest {
32            ip: ip.to_string(),
33            port,
34            counter,
35            protocol,
36            payload,
37            concurrency,
38        }
39    }
40
41    pub fn run(&self) -> io::Result<()> {
42        let tasks = self.counter / self.concurrency;
43        let tasks_remainder = self.counter % self.concurrency;
44        let ip = Arc::new(self.ip.clone());
45        let payload = Arc::new(self.payload.clone());
46
47        let mut threads = vec![];
48
49        println!("Starting test with {} threads", self.concurrency);
50
51        for _ in 0..self.concurrency {
52            let ip = Arc::clone(&ip);
53            let payload = Arc::clone(&payload);
54            let pentest = Pentest {
55                ip: ip.to_string(),
56                port: self.port,
57                counter: tasks,
58                protocol: self.protocol.clone(),
59                payload: payload.to_vec(),
60                concurrency: 1,
61            };
62
63            threads.push(thread::spawn(move || {
64                pentest.run_inner().unwrap();
65            }));
66        }
67
68        if tasks_remainder > 0 {
69            let pentest = Pentest {
70                ip: self.ip.clone(),
71                port: self.port,
72                counter: tasks_remainder,
73                protocol: self.protocol.clone(),
74                payload: self.payload.clone(),
75                concurrency: 1,
76            };
77
78            pentest.run_inner()?;
79        }
80
81        for (i, thread) in threads.into_iter().enumerate() {
82            println!("Joining thread {}", i + 1);
83            match thread.join() {
84                Ok(_) => println!("Thread {} joined", i + 1),
85                Err(err) => eprintln!("Error in thread {}: {:?}", i + 1, err),
86            }
87        }
88
89        Ok(())
90    }
91
92    fn run_inner(&self) -> io::Result<()> {
93        match self.protocol {
94            Protocol::TCP => self.run_tcp(),
95            Protocol::UDP => self.run_udp(),
96        }
97    }
98
99    fn run_tcp(&self) -> io::Result<()> {
100        let addr = format!("{}:{}", self.ip, self.port);
101
102        for i in 0..self.counter {
103            println!("Sending TCP packet {} to {}:{}", i + 1, self.ip, self.port);
104            match TcpStream::connect(&addr) {
105                Ok(mut stream) => {
106                    stream.set_write_timeout(Some(Duration::from_secs(1)))?;
107                    stream.write_all(&self.payload)?;
108                }
109                Err(err) => eprintln!("Failed to send TCP packet {}: {}", i + 1, err),
110            }
111        }
112
113        println!("Finished sending TCP packets");
114
115        Ok(())
116    }
117
118    fn run_udp(&self) -> io::Result<()> {
119        let socket = UdpSocket::bind("0.0.0.0:0")?;
120        socket.set_write_timeout(Some(Duration::from_secs(1)))?;
121
122        let addr = format!("{}:{}", self.ip, self.port);
123
124        for i in 0..self.counter {
125            println!("Sending UDP packet {} to {}:{}", i + 1, self.ip, self.port);
126            if let Err(err) = socket.send_to(&self.payload, &addr) {
127                eprintln!("Failed to send UDP packet {}: {}", i + 1, err);
128            }
129        }
130
131        println!("Finished sending UDP packets");
132
133        Ok(())
134    }
135}
pentest/lib.rs

pentest/
lib.rs
