protest_extras/generators/
network.rs1use protest::{Generator, GeneratorConfig};
11use rand::Rng;
12use std::net::{Ipv4Addr, Ipv6Addr};
13
14#[derive(Debug, Clone)]
20pub enum IpAddressGenerator {
21 V4,
23 V6,
25 Both,
27}
28
29impl IpAddressGenerator {
30 pub fn ipv4() -> Self {
32 Self::V4
33 }
34
35 pub fn ipv6() -> Self {
37 Self::V6
38 }
39
40 pub fn both() -> Self {
42 Self::Both
43 }
44}
45
46impl Generator<String> for IpAddressGenerator {
47 fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> String {
48 match self {
49 IpAddressGenerator::V4 => {
50 let octets: [u8; 4] = [rng.r#gen(), rng.r#gen(), rng.r#gen(), rng.r#gen()];
51 Ipv4Addr::from(octets).to_string()
52 }
53 IpAddressGenerator::V6 => {
54 let segments: [u16; 8] = [
55 rng.r#gen(),
56 rng.r#gen(),
57 rng.r#gen(),
58 rng.r#gen(),
59 rng.r#gen(),
60 rng.r#gen(),
61 rng.r#gen(),
62 rng.r#gen(),
63 ];
64 Ipv6Addr::from(segments).to_string()
65 }
66 IpAddressGenerator::Both => {
67 if rng.r#gen::<bool>() {
68 IpAddressGenerator::V4.generate(rng, _config)
69 } else {
70 IpAddressGenerator::V6.generate(rng, _config)
71 }
72 }
73 }
74 }
75
76 fn shrink(&self, value: &String) -> Box<dyn Iterator<Item = String>> {
77 let mut shrinks = Vec::new();
79
80 if value != "127.0.0.1" && value != "::1" {
82 shrinks.push("127.0.0.1".to_string());
83 }
84
85 if value != "0.0.0.0" && value != "::" {
87 shrinks.push("0.0.0.0".to_string());
88 }
89
90 Box::new(shrinks.into_iter())
91 }
92}
93
94#[derive(Debug, Clone)]
100pub struct EmailGenerator {
101 min_local_len: usize,
102 max_local_len: usize,
103 min_domain_len: usize,
104 max_domain_len: usize,
105}
106
107impl EmailGenerator {
108 pub fn new() -> Self {
110 Self {
111 min_local_len: 3,
112 max_local_len: 20,
113 min_domain_len: 5,
114 max_domain_len: 30,
115 }
116 }
117
118 pub fn with_lengths(
120 min_local_len: usize,
121 max_local_len: usize,
122 min_domain_len: usize,
123 max_domain_len: usize,
124 ) -> Self {
125 Self {
126 min_local_len,
127 max_local_len,
128 min_domain_len,
129 max_domain_len,
130 }
131 }
132
133 fn generate_local_part(&self, rng: &mut dyn rand::RngCore) -> String {
134 let len = rng.r#gen_range(self.min_local_len..=self.max_local_len);
135 let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789._-";
136
137 (0..len)
138 .map(|i| {
139 if i == 0 {
140 let alphanumeric = b"abcdefghijklmnopqrstuvwxyz0123456789";
142 alphanumeric[rng.r#gen_range(0..alphanumeric.len())] as char
143 } else {
144 valid_chars[rng.r#gen_range(0..valid_chars.len())] as char
145 }
146 })
147 .collect()
148 }
149
150 fn generate_domain(&self, rng: &mut dyn rand::RngCore) -> String {
151 let len = rng.r#gen_range(self.min_domain_len..=self.max_domain_len);
152 let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789-";
153
154 let domain_name: String = (0..len)
155 .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
156 .collect();
157
158 let tlds = ["com", "org", "net", "edu", "gov", "io", "dev"];
160 let tld = tlds[rng.r#gen_range(0..tlds.len())];
161
162 format!("{}.{}", domain_name, tld)
163 }
164}
165
166impl Default for EmailGenerator {
167 fn default() -> Self {
168 Self::new()
169 }
170}
171
172impl Generator<String> for EmailGenerator {
173 fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> String {
174 let local = self.generate_local_part(rng);
175 let domain = self.generate_domain(rng);
176 format!("{}@{}", local, domain)
177 }
178
179 fn shrink(&self, value: &String) -> Box<dyn Iterator<Item = String>> {
180 let mut shrinks = Vec::new();
181
182 if value != "a@b.com" {
184 shrinks.push("a@b.com".to_string());
185 }
186
187 if let Some((local, domain)) = value.split_once('@') {
189 if local.len() > 1 {
190 shrinks.push(format!("{}@{}", &local[..local.len() - 1], domain));
191 }
192
193 if let Some((domain_name, tld)) = domain.rsplit_once('.')
195 && domain_name.len() > 1
196 {
197 shrinks.push(format!(
198 "{}@{}.{}",
199 local,
200 &domain_name[..domain_name.len() - 1],
201 tld
202 ));
203 }
204 }
205
206 Box::new(shrinks.into_iter())
207 }
208}
209
210#[derive(Debug, Clone)]
216pub struct UrlGenerator {
217 schemes: Vec<&'static str>,
218 include_port: bool,
219 include_path: bool,
220 include_query: bool,
221}
222
223impl UrlGenerator {
224 pub fn new() -> Self {
226 Self {
227 schemes: vec!["http", "https"],
228 include_port: true,
229 include_path: true,
230 include_query: true,
231 }
232 }
233
234 pub fn with_schemes(schemes: Vec<&'static str>) -> Self {
236 Self {
237 schemes,
238 include_port: true,
239 include_path: true,
240 include_query: true,
241 }
242 }
243
244 pub fn with_port(mut self, include: bool) -> Self {
246 self.include_port = include;
247 self
248 }
249
250 pub fn with_path(mut self, include: bool) -> Self {
252 self.include_path = include;
253 self
254 }
255
256 pub fn with_query(mut self, include: bool) -> Self {
258 self.include_query = include;
259 self
260 }
261
262 fn generate_host(&self, rng: &mut dyn rand::RngCore) -> String {
263 let len = rng.r#gen_range(5..20);
264 let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789-";
265
266 let domain: String = (0..len)
267 .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
268 .collect();
269
270 let tlds = ["com", "org", "net", "io", "dev"];
271 let tld = tlds[rng.r#gen_range(0..tlds.len())];
272
273 format!("{}.{}", domain, tld)
274 }
275
276 fn generate_path(&self, rng: &mut dyn rand::RngCore) -> String {
277 let num_segments = rng.r#gen_range(1..=4);
278 let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789_-";
279
280 let segments: Vec<String> = (0..num_segments)
281 .map(|_| {
282 let len = rng.r#gen_range(3..10);
283 (0..len)
284 .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
285 .collect()
286 })
287 .collect();
288
289 format!("/{}", segments.join("/"))
290 }
291
292 fn generate_query(&self, rng: &mut dyn rand::RngCore) -> String {
293 let num_params = rng.r#gen_range(1..=3);
294 let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789";
295
296 let params: Vec<String> = (0..num_params)
297 .map(|_| {
298 let key_len = rng.r#gen_range(3..8);
299 let val_len = rng.r#gen_range(3..10);
300
301 let key: String = (0..key_len)
302 .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
303 .collect();
304
305 let val: String = (0..val_len)
306 .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
307 .collect();
308
309 format!("{}={}", key, val)
310 })
311 .collect();
312
313 format!("?{}", params.join("&"))
314 }
315}
316
317impl Default for UrlGenerator {
318 fn default() -> Self {
319 Self::new()
320 }
321}
322
323impl Generator<String> for UrlGenerator {
324 fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> String {
325 let scheme = self.schemes[rng.r#gen_range(0..self.schemes.len())];
326 let host = self.generate_host(rng);
327
328 let mut url = format!("{}://{}", scheme, host);
329
330 if self.include_port && rng.r#gen::<bool>() {
331 let port: u16 = rng.r#gen_range(1024..65535);
332 url.push_str(&format!(":{}", port));
333 }
334
335 if self.include_path && rng.r#gen::<bool>() {
336 url.push_str(&self.generate_path(rng));
337 }
338
339 if self.include_query && rng.r#gen::<bool>() {
340 url.push_str(&self.generate_query(rng));
341 }
342
343 url
344 }
345
346 fn shrink(&self, value: &String) -> Box<dyn Iterator<Item = String>> {
347 let mut shrinks = Vec::new();
348
349 if value != "http://a.com" {
351 shrinks.push("http://a.com".to_string());
352 }
353
354 if let Some(base) = value.split('?').next()
356 && base != value
357 {
358 shrinks.push(base.to_string());
359 }
360
361 if let Some(scheme_and_host) = value
363 .split('/')
364 .take(3)
365 .collect::<Vec<_>>()
366 .join("/")
367 .into()
368 && scheme_and_host != *value
369 {
370 shrinks.push(scheme_and_host);
371 }
372
373 Box::new(shrinks.into_iter())
374 }
375}
376
377#[cfg(test)]
378mod tests {
379 use super::*;
380 use rand::thread_rng;
381
382 #[test]
383 fn test_ipv4_generator() {
384 let generator = IpAddressGenerator::ipv4();
385 let mut rng = thread_rng();
386 let config = GeneratorConfig::default();
387
388 for _ in 0..10 {
389 let ip = generator.generate(&mut rng, &config);
390 let parts: Vec<&str> = ip.split('.').collect();
392 assert_eq!(parts.len(), 4);
393 for part in parts {
394 let _num: u8 = part.parse().unwrap();
396 }
397 }
398 }
399
400 #[test]
401 fn test_email_generator() {
402 let generator = EmailGenerator::new();
403 let mut rng = thread_rng();
404 let config = GeneratorConfig::default();
405
406 for _ in 0..10 {
407 let email = generator.generate(&mut rng, &config);
408 assert!(email.contains('@'));
409 assert_eq!(email.matches('@').count(), 1);
410
411 let parts: Vec<&str> = email.split('@').collect();
412 assert_eq!(parts.len(), 2);
413 assert!(!parts[0].is_empty());
414 assert!(!parts[1].is_empty());
415 assert!(parts[1].contains('.'));
416 }
417 }
418
419 #[test]
420 fn test_url_generator() {
421 let generator = UrlGenerator::new();
422 let mut rng = thread_rng();
423 let config = GeneratorConfig::default();
424
425 for _ in 0..10 {
426 let url = generator.generate(&mut rng, &config);
427 assert!(url.starts_with("http://") || url.starts_with("https://"));
428 }
429 }
430}