zeloxy 0.5.2

use std::sync::Arc;
use std::time::Duration;

use tokio::net::TcpStream;

#[cfg(feature = "https")]
use rustls::pki_types::ServerName;
#[cfg(feature = "https")]
use rustls::{ClientConfig, RootCertStore};
#[cfg(feature = "https")]
use tokio_rustls::TlsConnector;
#[cfg(feature = "https")]
use tokio_rustls::client::TlsStream;
#[cfg(feature = "https")]
use webpki_roots::TLS_SERVER_ROOTS;

#[cfg(feature = "http")]
use crate::proxy::protocols::connect_http;
#[cfg(feature = "https")]
use crate::proxy::protocols::connect_https;
#[cfg(feature = "socks4")]
use crate::proxy::protocols::connect_socks4;
#[cfg(feature = "socks5")]
use crate::proxy::protocols::connect_socks5;

use crate::proxy::{ProxyAuth, ProxyConnection, validate_proxy_str};
use crate::{ErrorKind, ProxyError, ProxyResult};

/// Структура прокси.
///
/// Поддерживаемые протоколы прокси:
///
/// - **HTTP**
/// - **HTTPS**
/// - **SOCKS4**
/// - **SOCKS5**
///
/// ## Примеры
///
/// ```rust, ignore
/// use tokio::io::{AsyncReadExt, AsyncWriteExt};
/// use zeloxy::important::{Proxy, ProxyProtocol};
///
/// #[tokio::main]
/// async fn main() -> std::io::Result<()> {
///   // Создаём HTTP-прокси и задаём адрес целевого сервера
///   let proxy = Proxy::new("91.132.92.231:80", ProxyProtocol::Http);
///
///   match proxy.connect("example.com", 80).await {
///     Ok(mut conn) => {
///       // Отправляем GET-запрос
///       conn.write_all(b"GET / HTTP/1.0\r\nHost: example.com\r\n\r\n").await?;
///
///       // Читаем ответ
///       let mut resp = Vec::new();
///       conn.read_to_end(&mut resp).await?;
///     
///       // Логгируем ответ
///       println!("{}", String::from_utf8_lossy(&resp));
///     },
///     Err(_) => {}, // Просто игнорируем ошибки
///   }
///
///   Ok(())
/// }
/// ```
///
/// Больше актуальных примеров: [смотреть](https://codeberg.org/nullclyze/zeloxy/src/branch/main/examples)
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Proxy {
  addr: String,
  protocol: ProxyProtocol,
  timeout: u64,
  auth: Option<ProxyAuth>,
}

/// Функция получения прокси по умолчанию опираясь на включенные фичи
#[allow(unused)]
fn default_proxy() -> Proxy {
  #[cfg(feature = "http")]
  return Proxy::new("127.0.0.1:80", ProxyProtocol::Http);
  #[cfg(feature = "https")]
  return Proxy::new("127.0.0.1:443", ProxyProtocol::Https);
  #[cfg(feature = "socks4")]
  return Proxy::new("127.0.0.1:4145", ProxyProtocol::Socks4);
  #[cfg(feature = "socks5")]
  return Proxy::new("127.0.0.1:1080", ProxyProtocol::Socks5);
}

/// Протокол прокси
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProxyProtocol {
  #[cfg(feature = "http")]
  Http,

  #[cfg(feature = "https")]
  Https,

  #[cfg(feature = "socks4")]
  Socks4,

  #[cfg(feature = "socks5")]
  Socks5,
}

impl ProxyProtocol {
  /// Метод получения порта по умолчанию для текущего протокола
  pub fn default_port(&self) -> u16 {
    match self {
      #[cfg(feature = "http")]
      Self::Http => 80,
      #[cfg(feature = "https")]
      Self::Https => 443,
      #[cfg(feature = "socks4")]
      Self::Socks4 => 4145,
      #[cfg(feature = "socks5")]
      Self::Socks5 => 1080,
    }
  }
}

impl From<&str> for ProxyProtocol {
  fn from(value: &str) -> Self {
    match value {
      #[cfg(feature = "http")]
      "http" => Self::Http,
      #[cfg(feature = "https")]
      "https" => Self::Https,
      #[cfg(feature = "socks4")]
      "socks4" => Self::Socks4,
      #[cfg(feature = "socks5")]
      "socks5" => Self::Socks5,
      _ => Self::Socks5,
    }
  }
}

impl From<String> for ProxyProtocol {
  fn from(value: String) -> Self {
    Self::from(value.as_str())
  }
}

impl Proxy {
  /// Метод создания нового прокси
  pub fn new(addr: impl Into<String>, protocol: impl Into<ProxyProtocol>) -> Self {
    Self {
      addr: addr.into(),
      protocol: protocol.into(),
      timeout: 20000,
      auth: None,
    }
  }

  /// Метод создания нового прокси с авторизацией
  pub fn new_with_auth(addr: impl Into<String>, protocol: impl Into<ProxyProtocol>, auth: ProxyAuth) -> Self {
    Self {
      addr: addr.into(),
      protocol: protocol.into(),
      timeout: 20000,
      auth: Some(auth.into()),
    }
  }

  /// Метод установки авторизации прокси
  pub fn with_auth(mut self, auth: impl Into<ProxyAuth>) -> Self {
    self.auth = Some(auth.into());
    self
  }

  /// Метод установки таймаута подключения к прокси
  pub fn with_timeout(mut self, timeout: u64) -> Self {
    self.timeout = timeout;
    self
  }

  /// Метод установки протокола прокси
  pub fn with_protocol(mut self, protocol: impl Into<ProxyProtocol>) -> Self {
    self.protocol = protocol.into();
    self
  }

  /// Метод попытки создания соединения с прокси
  pub async fn is_available(&self) -> bool {
    match tokio::time::timeout(Duration::from_millis(self.timeout), TcpStream::connect(&self.addr)).await {
      Ok(result) => match result {
        Ok(_) => return true,
        Err(_) => return false,
      },
      Err(_) => return false,
    }
  }

  /// Метод получения протокола прокси
  pub fn protocol(&self) -> &ProxyProtocol {
    &self.protocol
  }

  /// Метод получения IP прокси
  pub fn ip(&self) -> String {
    self.addr.split(":").collect::<Vec<&str>>()[0].to_string()
  }

  /// Метод получения порта прокси
  pub fn port(&self) -> u16 {
    let port_str = self.addr.split(":").collect::<Vec<&str>>()[1];
    port_str.parse::<u16>().unwrap_or(self.protocol.default_port())
  }

  /// Метод получения адреса прокси в формате `IP:PORT`
  pub fn addr(&self) -> &str {
    &self.addr
  }

  /// Метод получения полного адреса прокси в формате `PROTOCOL://IP:PORT`
  pub fn full_addr(&self) -> String {
    let protocol = match self.protocol {
      #[cfg(feature = "http")]
      ProxyProtocol::Http => "http",
      #[cfg(feature = "https")]
      ProxyProtocol::Https => "https",
      #[cfg(feature = "socks4")]
      ProxyProtocol::Socks4 => "socks4",
      #[cfg(feature = "socks5")]
      ProxyProtocol::Socks5 => "socks5",
    };

    format!("{}://{}", protocol, self.addr)
  }

  /// Метод установки TCP-соединения
  async fn connect_tcp(&self) -> ProxyResult<TcpStream> {
    match tokio::time::timeout(Duration::from_millis(self.timeout), TcpStream::connect(&self.addr)).await {
      Ok(result) => match result {
        Ok(s) => Ok(s),
        Err(_) => Err(ProxyError::new(ErrorKind::NotConnected, "could not connect to specified server")),
      },
      Err(_) => Err(ProxyError::new(
        ErrorKind::Timeout,
        "failed to connect to server within specified time",
      )),
    }
  }

  /// Функция установки TLS-соединения
  #[cfg(feature = "https")]
  async fn connect_tls(&self) -> ProxyResult<TlsStream<TcpStream>> {
    let tcp_stream = self.connect_tcp().await?;
    let server_name = if let Ok(ip) = self.ip().parse::<std::net::IpAddr>() {
      ServerName::IpAddress(ip.into())
    } else {
      ServerName::try_from(self.ip()).map_err(|_| ProxyError::new(ErrorKind::InvalidData, "invalid proxy server name"))?
    };

    let mut root_store = RootCertStore::empty();
    root_store.extend(TLS_SERVER_ROOTS.iter().cloned());

    let config = ClientConfig::builder().with_root_certificates(root_store).with_no_client_auth();

    let connector = TlsConnector::from(Arc::new(config));

    connector
      .connect(server_name, tcp_stream)
      .await
      .map_err(|e| ProxyError::new(ErrorKind::NotConnected, format!("TLS handshake with proxy failed: {}", e)))
  }

  /// Метод подключения к прокси
  pub async fn connect(&self, target_host: impl Into<String>, target_port: u16) -> ProxyResult<ProxyConnection> {
    let target_host = target_host.into();

    match self.protocol {
      #[cfg(feature = "https")]
      ProxyProtocol::Https => {
        let mut stream = self.connect_tls().await?;
        connect_https(&mut stream, target_host, target_port, &self.auth).await?;
        Ok(ProxyConnection::Tls(stream))
      }
      _ => {
        let mut stream = self.connect_tcp().await?;

        match self.protocol {
          #[cfg(feature = "http")]
          ProxyProtocol::Http => connect_http(&mut stream, target_host, target_port, &self.auth).await?,
          #[cfg(feature = "socks5")]
          ProxyProtocol::Socks5 => connect_socks5(&mut stream, target_host, target_port, &self.auth).await?,
          #[cfg(feature = "socks4")]
          ProxyProtocol::Socks4 => connect_socks4(&mut stream, target_host, target_port, &self.auth).await?,
          #[allow(unreachable_patterns)]
          _ => {}
        }

        Ok(ProxyConnection::Tcp(stream))
      }
    }
  }

  /// Метод подключения к прокси с ранее созданным TCP-соединением
  pub async fn connect_with_stream(&self, stream: &mut TcpStream, target_host: impl Into<String>, target_port: u16) -> ProxyResult<()> {
    match self.protocol {
      #[cfg(feature = "http")]
      ProxyProtocol::Http => connect_http(stream, target_host.into(), target_port, &self.auth).await?,
      #[cfg(feature = "https")]
      ProxyProtocol::Https => {
        return Err(ProxyError::new(
          ErrorKind::Unsupported,
          "this method does not support the HTTPS protocol",
        ));
      }
      #[cfg(feature = "socks4")]
      ProxyProtocol::Socks4 => connect_socks4(stream, target_host.into(), target_port, &self.auth).await?,
      #[cfg(feature = "socks5")]
      ProxyProtocol::Socks5 => connect_socks5(stream, target_host.into(), target_port, &self.auth).await?,
    }

    Ok(())
  }
}

impl From<String> for Proxy {
  fn from(value: String) -> Self {
    if !validate_proxy_str(&value) {
      return default_proxy();
    }

    let proxy_split = value.split("://").collect::<Vec<&str>>();
    let without_protocol = proxy_split[1].split('@').collect::<Vec<&str>>();

    let (protocol, basic_auth, addr) = {
      if without_protocol.len() == 2 {
        (proxy_split[0], Some(without_protocol[0]), without_protocol[1])
      } else {
        (proxy_split[0], None, without_protocol[0])
      }
    };

    if let Some(auth) = basic_auth {
      let auth_split = auth.split(':').collect::<Vec<&str>>();

      if auth_split.len() != 2 {
        Self::new(addr, protocol)
      } else {
        Self::new_with_auth(
          addr,
          protocol,
          ProxyAuth::Basic {
            username: auth_split[0].to_string(),
            password: auth_split[1].to_string(),
          },
        )
      }
    } else {
      Self::new(addr, protocol)
    }
  }
}

impl From<&str> for Proxy {
  fn from(value: &str) -> Self {
    if !validate_proxy_str(value) {
      return default_proxy();
    }

    let proxy_split = value.split("://").collect::<Vec<&str>>();
    let without_protocol = proxy_split[1].split('@').collect::<Vec<&str>>();

    let (protocol, basic_auth, addr) = {
      if without_protocol.len() == 2 {
        (proxy_split[0], Some(without_protocol[0]), without_protocol[1])
      } else {
        (proxy_split[0], None, without_protocol[0])
      }
    };

    if let Some(auth) = basic_auth {
      let auth_split = auth.split(':').collect::<Vec<&str>>();

      if auth_split.len() != 2 {
        Self::new(addr, protocol)
      } else {
        Self::new_with_auth(
          addr,
          protocol,
          ProxyAuth::Basic {
            username: auth_split[0].to_string(),
            password: auth_split[1].to_string(),
          },
        )
      }
    } else {
      Self::new(addr, protocol)
    }
  }
}

impl From<Arc<Proxy>> for Proxy {
  fn from(value: Arc<Proxy>) -> Self {
    Self {
      addr: value.addr.clone(),
      protocol: value.protocol.clone(),
      timeout: value.timeout,
      auth: value.auth.clone(),
    }
  }
}