use crate::client_response::ClientResponse;
use crate::stream::ClientStream;
use crate::url::Url;
use crate::{Method, Protocol};
use json::JsonValue;
use log::debug;
use rand::distr::Alphanumeric;
use rand::RngExt;
use rustls::RootCertStore;
use rustls::{ClientConnection, StreamOwned};
use std::collections::HashMap;
use std::io::Read;
use std::net::{IpAddr, TcpStream};
use std::path::Path;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use std::{fs, thread};
const REDIRECT_STATUS_CODES: [u16; 3] = [301, 302, 303];
pub struct Client {
debug: bool,
url: String,
protocol: Protocol,
method: Method,
query: HashMap<String, String>,
headers: HashMap<String, JsonValue>,
cookies: HashMap<String, String>,
body: Vec<u8>,
read_timeout: u64,
write_timeout: u64,
auto_location: bool,
boundary: String,
retry: usize,
ca_pem: String,
pkcs12: String,
pkcs12_password: String,
}
impl Default for Client {
fn default() -> Self {
Self::new()
}
}
impl Client {
pub fn new() -> Self {
Self {
debug: false,
url: "".to_string(),
protocol: Protocol::HTTP1_1,
method: Method::Other("".to_string()),
query: HashMap::new(),
headers: Default::default(),
cookies: Default::default(),
body: vec![],
read_timeout: 10,
write_timeout: 10,
auto_location: false,
boundary: "".to_string(),
retry: 0,
ca_pem: "".to_string(),
pkcs12: "".to_string(),
pkcs12_password: "".to_string(),
}
}
pub fn reset(&mut self) -> &mut Self {
self.debug = false;
self.url = "".to_string();
self.protocol = Protocol::HTTP1_1;
self.method = Method::Other("".to_string());
self.query = HashMap::new();
self.headers = Default::default();
self.cookies = Default::default();
self.body = vec![];
self.read_timeout = 0;
self.write_timeout = 0;
self.auto_location = false;
self.boundary = "".to_string();
self.retry = 0;
self
}
pub fn debug(&mut self) -> &mut Self {
self.debug = true;
self
}
pub fn ca_pem(&mut self, ca_pem: &str) -> &mut Self {
self.ca_pem = ca_pem.to_string();
self
}
pub fn pkcs12(&mut self, pkcs12: &str, password: &str) -> &mut Self {
self.pkcs12 = pkcs12.to_string();
self.pkcs12_password = password.to_string();
self
}
pub fn write_timeout(&mut self, time: u64) -> &mut Self {
self.write_timeout = time;
self
}
pub fn read_timeout(&mut self, time: u64) -> &mut Self {
self.read_timeout = time;
self
}
pub fn http10(&mut self) -> &mut Self {
self.protocol = Protocol::HTTP1_0;
self
}
pub fn http11(&mut self) -> &mut Self {
self.protocol = Protocol::HTTP1_1;
self
}
pub fn http2(&mut self) -> &mut Self {
self.protocol = Protocol::HTTP2;
self
}
pub fn retry(&mut self, count: usize) -> &mut Self {
self.retry = count;
self
}
pub fn auto_location(&mut self) -> &mut Self {
self.auto_location = true;
self
}
fn generate_boundary(&mut self) -> String {
let rand: String = rand::rng()
.sample_iter(&Alphanumeric)
.take(16)
.map(char::from)
.collect();
format!("----WebKitFormBoundary{}", rand)
}
pub fn get(&mut self, url: &str) -> &mut Self {
self.url = url.to_string();
self.method = Method::GET;
self
}
pub fn post(&mut self, url: &str) -> &mut Self {
self.url = url.to_string();
self.method = Method::POST;
self
}
pub fn head(&mut self, url: &str) -> &mut Self {
self.url = url.to_string();
self.method = Method::HEAD;
self
}
pub fn delete(&mut self, url: &str) -> &mut Self {
self.url = url.to_string();
self.method = Method::DELETE;
self
}
pub fn method(&mut self, method: &str, url: &str) -> &mut Self {
self.url = url.to_string();
self.method = Method::Other(method.to_string());
self
}
pub fn header(&mut self, key: &str, value: JsonValue) -> &mut Self {
self.headers.insert(key.to_string(), value);
self
}
pub fn query(&mut self, key: &str, value: JsonValue) -> &mut Self {
let value = br_crypto::encoding::urlencoding_encode(value.to_string().as_str());
self.query.insert(key.to_string(), value.to_string());
self
}
pub fn raw_json(&mut self, data: JsonValue) -> &mut Self {
self.header("Content-Type", "application/json".into());
self.body = data.to_string().as_bytes().to_vec();
self.header("Content-Length", self.body.len().into());
self
}
pub fn form_data(&mut self, data: JsonValue) -> &mut Self {
self.boundary = self.generate_boundary();
self.header(
"Content-Type",
format!("multipart/form-data; boundary={}", self.boundary).into(),
);
let mut body = vec![];
for (key, value) in data.entries() {
let file = Path::new(value.as_str().unwrap_or(""));
if file.is_file() {
let mut t = vec![];
let r = match fs::File::open(file).and_then(|mut f| f.read(&mut t)) {
Ok(r) => r,
Err(_) => continue,
};
let filename = file
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("unknown")
.to_string();
let content_type = if let Some(kind) = infer::get(&t[..r]) {
kind.mime_type()
} else {
"application/octet-stream"
};
body.extend(format!("--{}\r\n", self.boundary).as_bytes().to_vec());
body.extend(
format!(
r#"Content-Disposition: form-data; name="{}"; filename="{}"{}"#,
key, filename, "\r\n"
)
.as_bytes()
.to_vec(),
);
body.extend(
format!(r#"Content-Type: {}{}"#, content_type, "\r\n\r\n")
.as_bytes()
.to_vec(),
);
body.extend(&t[..r]);
body.extend("\r\n".to_string().as_bytes().to_vec());
} else {
body.extend(format!("--{}\r\n", self.boundary).as_bytes().to_vec());
body.extend(
format!(
r#"Content-Disposition: form-data; name="{}"{}"#,
key, "\r\n\r\n"
)
.as_bytes()
.to_vec(),
);
body.extend(format!("{value}\r\n").as_bytes().to_vec());
}
}
body.extend(format!("--{}--\r\n", self.boundary).as_bytes().to_vec());
self.body = body.to_vec();
self.header("Content-Length", self.body.len().into());
self
}
pub fn x_www_form_urlencoded(&mut self, data: JsonValue) -> &mut Self {
self.header("Content-Type", "application/x-www-form-urlencoded".into());
let mut params = vec![];
for (key, value) in data.entries() {
params.push(format!("{}={}", key, value));
}
let params = params.join("&");
self.body = params.as_bytes().to_vec();
self.header("Content-Length", self.body.len().into());
self
}
pub fn send(&mut self) -> Result<ClientResponse, String> {
let mut url = Url::parse(&self.url)?;
let stream = match TcpStream::connect(format!("{}:{}", url.host, url.port)) {
Ok(stream) => stream,
Err(e) => {
if self.retry > 0 {
self.retry -= 1;
return self.send();
}
return Err(e.to_string());
}
};
if self.read_timeout > 0 {
let _ = stream.set_read_timeout(Option::from(Duration::from_secs(self.read_timeout)));
}
if self.write_timeout > 0 {
let _ = stream.set_write_timeout(Option::from(Duration::from_secs(self.write_timeout)));
}
let mut stream = if url.scheme == "https" {
let root_store =
RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
let mut config = rustls::ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
config.alpn_protocols = vec![self.protocol.str().to_lowercase().as_bytes().to_vec()];
let example_com = url
.host
.clone()
.try_into()
.map_err(|e| format!("invalid DNS name: {}", e))?;
let client = ClientConnection::new(Arc::new(config), example_com)
.map_err(|e| format!("TLS connection error: {}", e))?;
ClientStream::Https(Box::new(StreamOwned::new(client, stream)))
} else {
ClientStream::Http(stream)
};
if !self.query.is_empty() {
url.set_query(self.query.clone());
}
let request_line = format!(
"{} {} {}",
self.method.str(),
url.path_query.as_str(),
self.protocol.str()
);
let host = if IpAddr::from_str(url.host.as_str()).is_ok() {
format!("Host: {}:{}", url.host, url.port)
} else {
format!("Host: {}", url.host)
};
let mut header = vec![request_line, host];
for (key, value) in self.headers.iter() {
header.push(format!("{}: {}", key, value));
}
for (key, value) in self.cookies.iter() {
header.push(format!("Set-Cookie: {}={}", key, value));
}
header.push("\r\n".to_string());
match stream.write_all(header.join("\r\n").as_bytes()) {
Ok(()) => {
match stream.flush() {
Ok(e) => e,
Err(e) => return Err(e.to_string()),
};
}
Err(e) => {
if self.retry > 0 {
self.retry -= 1;
return self.send();
}
return Err(e.to_string());
}
};
if self.debug {
debug!("\r\n==============请求头 {:?}============== \r\n{}\r\n============================",thread::current().id(), header.join("\r\n"));
}
if !self.body.is_empty() {
match stream.write_all(self.body.as_slice()) {
Ok(()) => {
match stream.flush() {
Ok(e) => e,
Err(e) => return Err(e.to_string()),
};
}
Err(e) => {
if self.retry > 0 {
self.retry -= 1;
return self.send();
}
return Err(e.to_string());
}
}
}
if self.debug && self.boundary.is_empty() {
debug!("\r\n==============请求体 {:?}============== \r\n{}\r\n============================",thread::current().id(),String::from_utf8_lossy(self.body.as_slice()));
}
let mut res = ClientResponse::new(self.debug, self.method.clone(), stream).handle()?;
if REDIRECT_STATUS_CODES.contains(&res.status()) && !res.location().is_empty() {
self.url = res.location();
return self.send();
}
self.reset();
Ok(res)
}
}