implement integrated http sward

ubw-sward/src/http/integrated.rs

@@ -0,0 +1,182 @@
+use crate::http::header_config::HeadersConfig;
+use crate::http::simple::HttpSwardArray;
+use crate::http::{RandomUrlGenerator, SimpleHttpSward};
+use crate::utils::multiplexed::MultiplexedSward;
+use rand::Rng;
+use reqwest::Method;
+use std::net::{IpAddr, SocketAddr};
+use std::sync::Arc;
+use thiserror::Error;
+use url::Url;
+use wyrand::WyRand;
+
+#[derive(Clone)]
+enum AttackTarget {
+    Random(RandomUrlGenerator<WyRand>),
+    Fixed(Url),
+}
+
+pub struct IntegratedHttpSward {
+    attack_target: AttackTarget,
+    request_sender: MultiplexedSward<HttpSwardArray<Box<[SimpleHttpSward]>>>,
+}
+
+impl IntegratedHttpSward {
+    pub fn builder() -> IntegratedHttpSwardBuilder {
+        IntegratedHttpSwardBuilder::new()
+    }
+}
+
+pub struct IntegratedHttpSwardBuilder {
+    capacity: usize,
+    headers_config: HeadersConfig,
+    method: Method,
+    ip_list: Box<[IpAddr]>,
+}
+
+impl Default for IntegratedHttpSwardBuilder {
+    fn default() -> Self {
+        Self {
+            capacity: 32,
+            method: Method::GET,
+            headers_config: HeadersConfig::default(),
+            ip_list: Box::new([]),
+        }
+    }
+}
+
+impl IntegratedHttpSwardBuilder {
+    pub fn new() -> Self {
+        Self::default()
+    }
+    pub fn set_capacity(mut self, cap: usize) -> Self {
+        self.capacity = cap;
+        self
+    }
+    pub fn set_method(mut self, method: Method) -> Self {
+        self.method = method;
+        self
+    }
+    pub fn set_headers_config(mut self, headers: HeadersConfig) -> Self {
+        self.headers_config = headers;
+        self
+    }
+    pub fn set_ip_list(mut self, ip: &[IpAddr]) -> Self {
+        let cloned: Box<[_]> = ip.iter().copied().collect();
+        self.ip_list = cloned;
+        self
+    }
+
+    fn build_clients(
+        self,
+        url_example: &Url,
+    ) -> Result<
+        impl Iterator<Item = Result<reqwest::Client, HttpSwardBuildError>>,
+        HttpSwardBuildError,
+    > {
+        // check the ip list
+        if self.ip_list.is_empty() {
+            return Err(HttpSwardBuildError::EmptyIpResolve);
+        }
+
+        // resolve the IP
+        let domain = url_example
+            .host_str()
+            .ok_or(HttpSwardBuildError::MissingHost)?;
+        let port = url_example.port_or_known_default().unwrap_or(443);
+        let socket_addrs = self
+            .ip_list
+            .into_iter()
+            .map(move |ip| SocketAddr::new(ip, port));
+
+        // get the cookie
+        let maybe_cookie = self
+            .headers_config
+            .get_cookie_jar(&url_example)
+            .map(Arc::new);
+
+        // build the clients
+        let clients = socket_addrs
+            .map(|addr| {
+                reqwest::Client::builder()
+                    .resolve(domain, addr)
+                    .use_native_tls()
+            })
+            .map(move |builder| self.headers_config.set_client_header(builder))
+            .map(move |builder| match &maybe_cookie {
+                Some(jar) => builder.cookie_provider(jar.clone()).cookie_store(true),
+                None => builder,
+            })
+            .map(|builder| {
+                builder
+                    .build()
+                    .map_err(HttpSwardBuildError::ReqwestBuildError)
+            });
+
+        Ok(clients)
+    }
+
+    pub fn url(self, url: Url) -> Result<IntegratedHttpSward, HttpSwardBuildError> {
+        let request_time_headers = self.headers_config.other_headers.clone();
+        let method = self.method.clone();
+        let capacity = self.capacity;
+        let clients = self.build_clients(&url)?;
+        let swards: Box<[_]> = clients
+            .map(|client_result| {
+                client_result.map(|client| {
+                    SimpleHttpSward::new(client, method.clone(), request_time_headers.clone())
+                })
+            })
+            .collect::<Result<_, HttpSwardBuildError>>()?;
+
+        Ok(IntegratedHttpSward {
+            attack_target: AttackTarget::Fixed(url),
+            request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity),
+        })
+    }
+
+    pub fn template(
+        self,
+        template: impl AsRef<str>,
+    ) -> Result<IntegratedHttpSward, HttpSwardBuildError> {
+        let request_time_headers = self.headers_config.other_headers.clone();
+        let method = self.method.clone();
+        let capacity = self.capacity;
+        let random_core = WyRand::new(rand::rng().random());
+
+        let mut random_url_generator = RandomUrlGenerator::new(random_core, template.as_ref())
+            .map_err(|_| HttpSwardBuildError::BadTemplate)?;
+        let example_url = random_url_generator
+            .generate_url()
+            .map_err(|_| HttpSwardBuildError::BadTemplate)?;
+
+        let clients = self.build_clients(&example_url)?;
+        let swards: Box<[_]> = clients
+            .map(|client_result| {
+                client_result.map(|client| {
+                    SimpleHttpSward::new(client, method.clone(), request_time_headers.clone())
+                })
+            })
+            .collect::<Result<_, HttpSwardBuildError>>()?;
+
+        Ok(IntegratedHttpSward {
+            attack_target: AttackTarget::Random(random_url_generator),
+            request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity),
+        })
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum HttpSwardBuildError {
+    #[error("Cannot parse domain or address in the url.")]
+    MissingHost,
+
+    #[error("The IP list is empty")]
+    EmptyIpResolve,
+
+    #[error("The template cannot generate a correct URL")]
+    BadTemplate,
+
+    #[error("Failed to build reqwest client {0}")]
+    ReqwestBuildError(#[from] reqwest::Error),
+}