ubw-network/ubw-sward/src/http/integrated.rs

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),
}