diff --git a/.gitignore b/.gitignore index 4ef4bff94..0cc6c6487 100644 --- a/.gitignore +++ b/.gitignore @@ -29,10 +29,4 @@ uploads ./vendor upgrade ca -chat3.openai.com.har -chat4.openai.com.har -platform.openai.com.har -login.chat.openai.com.har -4.har -auth0.openai.com_Archive.har -auth0.openai.com.har \ No newline at end of file +*.har \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index d71993d68..dddead4ae 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ anyhow = "1.0.75" clap = { version = "4.4.3", features = ["derive", "env"] } serde = {version = "1.0.188", features = ["derive"] } openai = { path = "openai" } +mitm ={ path = "mitm", optional = true } cidr = "0.2.2" toml = "0.8.0" url = "2.4.1" @@ -57,7 +58,8 @@ env_logger = "0.10.0" openai = { path = "openai" } [features] -default = ["serve"] +default = ["serve", "mitm"] +mitm = ["openai/preauth", "dep:mitm"] terminal = [ "openai/api", "dep:tokio", diff --git a/mitm/Cargo.toml b/mitm/Cargo.toml new file mode 100644 index 000000000..8e8b4dbfe --- /dev/null +++ b/mitm/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "mitm" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +log = "0.4.20" +anyhow = "1.0.75" +thiserror = "1.0.48" +async-trait = "0.1.73" +reqwest = { package = "reqwest-impersonate", version ="0.11.30", default-features = false, features = [ + "boring-tls", "impersonate", "stream", "socks" +] } +typed-builder = "0.18.0" +time = "0.3.30" +rand = "0.8.5" +moka = { version = "0.12.1", default-features = false, features = ["sync"] } +tokio = { version = "1.15.0", default-features = false } +rcgen = { version = "0.10", features = ["x509-parser"] } +hyper = { version = "0.14.27", default-features = false } +tokio-rustls = { version = "0.24.1", default-features = false, features = ["tls12"] } +rustls = { version = "0.21.8", features = ["dangerous_configuration"] } +wildmatch = "2.1" +http = "0.2.11" +pin-project = "1" +byteorder = "1.4" +rustls-pemfile = "1.0" \ No newline at end of file diff --git a/openai/src/serve/preauth/cagen.rs b/mitm/src/cagen.rs similarity index 90% rename from openai/src/serve/preauth/cagen.rs rename to mitm/src/cagen.rs index 0ed313bfb..f1409d793 100644 --- a/openai/src/serve/preauth/cagen.rs +++ b/mitm/src/cagen.rs @@ -1,9 +1,10 @@ +use log::error; use rcgen::Certificate; -use crate::{error, serve::preauth::proxy::CertificateAuthority}; - use std::fs; +use crate::proxy::CertificateAuthority; + pub fn gen_ca() -> Certificate { let cert = CertificateAuthority::gen_ca().expect("preauth generate cert"); let cert_crt = cert.serialize_pem().unwrap(); diff --git a/mitm/src/lib.rs b/mitm/src/lib.rs new file mode 100644 index 000000000..5014da921 --- /dev/null +++ b/mitm/src/lib.rs @@ -0,0 +1,60 @@ +pub mod cagen; +pub mod proxy; + +use anyhow::Context; +use std::{fs, net::SocketAddr, path::PathBuf}; +use typed_builder::TypedBuilder; + +use crate::proxy::{handler::HttpHandler, CertificateAuthority}; +use log::info; + +#[derive(TypedBuilder)] +pub struct Builder { + bind: SocketAddr, + upstream_proxy: Option, + cert: PathBuf, + key: PathBuf, + graceful_shutdown: tokio::sync::mpsc::Receiver<()>, + cerificate_cache_size: u32, + mitm_filters: Vec, + handler: T, +} + +impl Builder { + pub async fn mitm_proxy(self) -> anyhow::Result<()> { + info!("PreAuth CA Private key use: {}", self.key.display()); + let private_key_bytes = + fs::read(self.key).context("ca private key file path not valid!")?; + let private_key = rustls_pemfile::pkcs8_private_keys(&mut private_key_bytes.as_slice()) + .context("Failed to parse private key")?; + let key = rustls::PrivateKey(private_key[0].clone()); + + info!("PreAuth CA Certificate use: {}", self.cert.display()); + let ca_cert_bytes = fs::read(self.cert).context("ca cert file path not valid!")?; + let ca_cert = rustls_pemfile::certs(&mut ca_cert_bytes.as_slice()) + .context("Failed to parse CA certificate")?; + let cert = rustls::Certificate(ca_cert[0].clone()); + + let ca = CertificateAuthority::new( + key, + cert, + String::from_utf8(ca_cert_bytes).context("Failed to parse CA certificate")?, + self.cerificate_cache_size.into(), + ) + .context("Failed to create Certificate Authority")?; + + info!("PreAuth Http MITM Proxy listen on: http://{}", self.bind); + + let proxy = proxy::Proxy::builder() + .ca(ca.clone()) + .listen_addr(self.bind) + .upstream_proxy(self.upstream_proxy) + .mitm_filters(self.mitm_filters) + .handler(self.handler) + .graceful_shutdown(self.graceful_shutdown) + .build(); + + tokio::spawn(proxy.start_proxy()); + Ok(()) + } +} diff --git a/openai/src/serve/preauth/proxy/ca.rs b/mitm/src/proxy/ca.rs similarity index 100% rename from openai/src/serve/preauth/proxy/ca.rs rename to mitm/src/proxy/ca.rs diff --git a/openai/src/serve/preauth/proxy/error.rs b/mitm/src/proxy/error.rs similarity index 100% rename from openai/src/serve/preauth/proxy/error.rs rename to mitm/src/proxy/error.rs diff --git a/openai/src/serve/preauth/proxy/handler.rs b/mitm/src/proxy/handler.rs similarity index 100% rename from openai/src/serve/preauth/proxy/handler.rs rename to mitm/src/proxy/handler.rs diff --git a/openai/src/serve/preauth/proxy/http_client.rs b/mitm/src/proxy/http_client.rs similarity index 90% rename from openai/src/serve/preauth/proxy/http_client.rs rename to mitm/src/proxy/http_client.rs index af6a2833c..bfdda0bdf 100644 --- a/openai/src/serve/preauth/proxy/http_client.rs +++ b/mitm/src/proxy/http_client.rs @@ -1,5 +1,5 @@ use http::{response::Builder, Request, Response}; -use hyper::Body; +use hyper::{body, Body}; use reqwest::impersonate::Impersonate; use super::error::Error; @@ -33,7 +33,7 @@ impl HttpClient { .request(method, url) .headers(parts.headers) .version(parts.version) - .body(hyper::body::to_bytes(body).await?) + .body(body::to_bytes(body).await?) .send() .await?; @@ -46,6 +46,6 @@ impl HttpClient { .headers_mut() .map(|h| h.extend(resp.headers().clone())); - Ok(builder.body(hyper::body::Body::wrap_stream(resp.bytes_stream()))?) + Ok(builder.body(body::Body::wrap_stream(resp.bytes_stream()))?) } } diff --git a/openai/src/serve/preauth/proxy/mitm.rs b/mitm/src/proxy/mitm.rs similarity index 100% rename from openai/src/serve/preauth/proxy/mitm.rs rename to mitm/src/proxy/mitm.rs diff --git a/openai/src/serve/preauth/proxy/mod.rs b/mitm/src/proxy/mod.rs similarity index 100% rename from openai/src/serve/preauth/proxy/mod.rs rename to mitm/src/proxy/mod.rs diff --git a/openai/src/serve/preauth/proxy/sni_reader.rs b/mitm/src/proxy/sni_reader.rs similarity index 100% rename from openai/src/serve/preauth/proxy/sni_reader.rs rename to mitm/src/proxy/sni_reader.rs diff --git a/openai/Cargo.toml b/openai/Cargo.toml index 7a9ebfe56..605b54d88 100644 --- a/openai/Cargo.toml +++ b/openai/Cargo.toml @@ -43,6 +43,9 @@ nom = { version = "7.1.3", optional = true } mime = { version = "0.3.17", optional = true } futures-timer = { version = "3.0.2", optional = true } +# mitm +mitm = { path = "../mitm", optional = true } + # arkose aes = "0.8.3" md5 = "0.7.0" @@ -63,17 +66,6 @@ async-stream = { version = "0.3.5", optional = true } axum_csrf = { version = "0.7.2", features = ["layer"], optional = true } serde_urlencoded = { version = "0.7.1", optional = true } -# mitm -rcgen = { version = "0.10", features = ["x509-parser"], optional = true } -hyper = { version = "0.14.27", default-features = false, optional = true } -tokio-rustls = { version = "0.24.1", default-features = false, features = ["tls12"], optional = true } -rustls = { version = "0.21.8", features = ["dangerous_configuration"], optional = true } -wildmatch = { version = "2.1", optional = true } -http = { version = "0.2", optional = true } -pin-project = { version = "1", optional = true } -byteorder = { version = "1.4", optional = true } -rustls-pemfile = { version = "1.0", optional = true } - [target.'cfg(windows)'.dependencies.windows-sys] version = "0.48.0" default-features = false @@ -86,7 +78,7 @@ static-files = "0.2.3" default = ["serve", "limit", "template", "preauth"] api = ["stream"] serve = ["dep:serde_urlencoded", "dep:axum_csrf", "stream", "dep:async-stream", "dep:tracing", "dep:tracing-subscriber", "dep:tower-http", "dep:tower", "dep:bytes", "dep:time", "dep:axum-server", "dep:axum-extra", "dep:axum", "dep:static-files", "dep:futures-core", "dep:tera"] -preauth = ["dep:rustls-pemfile", "dep:rcgen", "dep:moka", "dep:hyper", "dep:tokio-rustls", "dep:rustls", "dep:wildmatch", "dep:http", "dep:pin-project", "dep:byteorder"] +preauth = ["dep:mitm"] stream = ["dep:tokio-util", "dep:futures", "dep:tokio-stream", "dep:eventsource-stream", "dep:futures-core", "dep:pin-project-lite", "dep:nom", "dep:mime", "dep:futures-timer"] remote-token = [] limit = ["dep:redis", "dep:redis-macros", "dep:moka"] diff --git a/openai/src/serve/middleware/csrf.rs b/openai/src/serve/middleware/csrf.rs index 0107b2239..d26270980 100644 --- a/openai/src/serve/middleware/csrf.rs +++ b/openai/src/serve/middleware/csrf.rs @@ -1,3 +1,4 @@ +use axum::http::{Method, Request, StatusCode}; use axum::{ body::{self, BoxBody, Full}, middleware::Next, @@ -5,7 +6,7 @@ use axum::{ Form, }; use axum_csrf::CsrfToken; -use http::{Method, Request, StatusCode}; +use mitm::proxy::hyper; use crate::{auth::model::AuthAccount, warn}; diff --git a/openai/src/serve/mod.rs b/openai/src/serve/mod.rs index 28a067450..82399686e 100644 --- a/openai/src/serve/mod.rs +++ b/openai/src/serve/mod.rs @@ -1,7 +1,7 @@ mod error; mod middleware; #[cfg(feature = "preauth")] -pub mod preauth; +mod preauth; mod proxy; mod puid; #[cfg(feature = "template")] @@ -14,11 +14,11 @@ use axum::body::Body; use axum::headers::authorization::Bearer; use axum::headers::Authorization; use axum::http::Response; +use axum::http::StatusCode; use axum::response::IntoResponse; use axum::routing::{any, get, post}; use axum::{Json, TypedHeader}; use axum_server::{AddrIncomingConfig, Handle}; -use http::StatusCode; use self::proxy::ext::RequestExt; use self::proxy::ext::SendRequestExt; @@ -189,16 +189,17 @@ impl Serve { // PreAuth mitm proxy #[cfg(feature = "preauth")] if let Some(pbind) = self.0.pbind.clone() { - if let Some(err) = preauth::mitm_proxy( - pbind, - self.0.pupstream.clone(), - self.0.pcert.clone(), - self.0.pkey.clone(), - rx, - ) - .await - .err() - { + let builder = mitm::Builder::builder() + .bind(pbind) + .upstream_proxy(self.0.pupstream.clone()) + .cert(self.0.pcert.clone()) + .key(self.0.pkey.clone()) + .graceful_shutdown(rx) + .cerificate_cache_size(1_000) + .mitm_filters(vec![String::from("ios.chat.openai.com")]) + .handler(preauth::PreAuthHanlder) + .build(); + if let Some(err) = builder.mitm_proxy().await.err() { crate::error!("PreAuth proxy error: {}", err); } } diff --git a/openai/src/serve/preauth.rs b/openai/src/serve/preauth.rs new file mode 100644 index 000000000..bb12f2619 --- /dev/null +++ b/openai/src/serve/preauth.rs @@ -0,0 +1,100 @@ +use mitm::proxy::hyper::{ + body::Body, + http::{header, HeaderMap, HeaderValue, Request, Response}, +}; +use mitm::proxy::{handler::HttpHandler, mitm::RequestOrResponse}; +use std::fmt::Write; + +use crate::{info, with_context}; + +#[derive(Clone)] +pub struct PreAuthHanlder; + +#[async_trait::async_trait] +impl HttpHandler for PreAuthHanlder { + async fn handle_request(&self, req: Request) -> RequestOrResponse { + if log::log_enabled!(log::Level::Debug) { + log_req(&req).await; + } + // extract preauth cookie + collect_preauth_cookie(req.headers()); + RequestOrResponse::Request(req) + } + + async fn handle_response(&self, res: Response) -> Response { + if log::log_enabled!(log::Level::Debug) { + log_res(&res).await; + } + collect_preauth_cookie(res.headers()); + res + } +} + +fn collect_preauth_cookie(headers: &HeaderMap) { + headers + .iter() + .filter(|(k, _)| k.eq(&header::COOKIE) || k.eq(&header::SET_COOKIE)) + .for_each(|(_, v)| { + let _ = v + .to_str() + .map(|value| with_context!(push_preauth_cookie, value)); + }); +} + +pub async fn log_req(req: &Request) { + let headers = req.headers(); + let mut header_formated = String::new(); + for (key, value) in headers { + let v = match value.to_str() { + Ok(v) => v.to_string(), + Err(_) => { + format!("[u8]; {}", value.len()) + } + }; + write!( + &mut header_formated, + "\t{:<20}{}\r\n", + format!("{}:", key.as_str()), + v + ) + .unwrap(); + } + + info!( + "{} {} +Headers: +{}", + req.method(), + req.uri().to_string(), + header_formated + ) +} + +pub async fn log_res(res: &Response) { + let headers = res.headers(); + let mut header_formated = String::new(); + for (key, value) in headers { + let v = match value.to_str() { + Ok(v) => v.to_string(), + Err(_) => { + format!("[u8]; {}", value.len()) + } + }; + write!( + &mut header_formated, + "\t{:<20}{}\r\n", + format!("{}:", key.as_str()), + v + ) + .unwrap(); + } + + info!( + "{} {:?} +Headers: +{}", + res.status(), + res.version(), + header_formated + ) +} diff --git a/openai/src/serve/preauth/mod.rs b/openai/src/serve/preauth/mod.rs deleted file mode 100644 index 96749e187..000000000 --- a/openai/src/serve/preauth/mod.rs +++ /dev/null @@ -1,153 +0,0 @@ -use std::{fs, net::SocketAddr, path::PathBuf}; - -use anyhow::Context; -use http::{header, Request, Response}; -use hyper::Body; - -use self::proxy::mitm::RequestOrResponse; -use crate::{ - serve::preauth::proxy::{handler::HttpHandler, CertificateAuthority}, - with_context, -}; -use log::info; - -pub mod cagen; -mod proxy; - -pub(super) async fn mitm_proxy( - bind: SocketAddr, - upstream_proxy: Option, - cert: PathBuf, - key: PathBuf, - graceful_shutdown: tokio::sync::mpsc::Receiver<()>, -) -> anyhow::Result<()> { - info!("PreAuth CA Private key use: {}", key.display()); - let private_key_bytes = fs::read(key).context("ca private key file path not valid!")?; - let private_key = rustls_pemfile::pkcs8_private_keys(&mut private_key_bytes.as_slice()) - .context("Failed to parse private key")?; - let key = rustls::PrivateKey(private_key[0].clone()); - - info!("PreAuth CA Certificate use: {}", cert.display()); - let ca_cert_bytes = fs::read(cert).context("ca cert file path not valid!")?; - let ca_cert = rustls_pemfile::certs(&mut ca_cert_bytes.as_slice()) - .context("Failed to parse CA certificate")?; - let cert = rustls::Certificate(ca_cert[0].clone()); - - let ca = CertificateAuthority::new( - key, - cert, - String::from_utf8(ca_cert_bytes).context("Failed to parse CA certificate")?, - 1_000, - ) - .context("Failed to create Certificate Authority")?; - - info!("PreAuth Http MITM Proxy listen on: http://{bind}"); - - let http_handler = PreAuthHanlder; - - let proxy = proxy::Proxy::builder() - .ca(ca.clone()) - .listen_addr(bind) - .upstream_proxy(upstream_proxy) - .mitm_filters(vec![String::from("ios.chat.openai.com")]) - .handler(http_handler.clone()) - .graceful_shutdown(graceful_shutdown) - .build(); - - tokio::spawn(proxy.start_proxy()); - Ok(()) -} - -#[derive(Clone)] -struct PreAuthHanlder; - -#[async_trait::async_trait] -impl HttpHandler for PreAuthHanlder { - async fn handle_request(&self, req: Request) -> RequestOrResponse { - if log::log_enabled!(log::Level::Info) { - log_req(&req).await; - } - // extract preauth cookie - collect_preauth_cookie(req.headers()); - RequestOrResponse::Request(req) - } - - async fn handle_response(&self, res: Response) -> Response { - if log::log_enabled!(log::Level::Info) { - log_res(&res).await; - } - collect_preauth_cookie(res.headers()); - res - } -} - -use std::fmt::Write; - -fn collect_preauth_cookie(headers: &http::HeaderMap) { - headers - .iter() - .filter(|(k, _)| k.eq(&header::COOKIE) || k.eq(&header::SET_COOKIE)) - .for_each(|(_, v)| { - let _ = v - .to_str() - .map(|value| with_context!(push_preauth_cookie, value)); - }); -} - -pub async fn log_req(req: &Request) { - let headers = req.headers(); - let mut header_formated = String::new(); - for (key, value) in headers { - let v = match value.to_str() { - Ok(v) => v.to_string(), - Err(_) => { - format!("[u8]; {}", value.len()) - } - }; - write!( - &mut header_formated, - "\t{:<20}{}\r\n", - format!("{}:", key.as_str()), - v - ) - .unwrap(); - } - - info!( - "{} {} -Headers: -{}", - req.method(), - req.uri().to_string(), - header_formated - ) -} - -pub async fn log_res(res: &Response) { - let headers = res.headers(); - let mut header_formated = String::new(); - for (key, value) in headers { - let v = match value.to_str() { - Ok(v) => v.to_string(), - Err(_) => { - format!("[u8]; {}", value.len()) - } - }; - write!( - &mut header_formated, - "\t{:<20}{}\r\n", - format!("{}:", key.as_str()), - v - ) - .unwrap(); - } - - info!( - "{} {:?} -Headers: -{}", - res.status(), - res.version(), - header_formated - ) -} diff --git a/openai/src/serve/proxy/toapi/mod.rs b/openai/src/serve/proxy/toapi/mod.rs index 285f072ec..7192e0d22 100644 --- a/openai/src/serve/proxy/toapi/mod.rs +++ b/openai/src/serve/proxy/toapi/mod.rs @@ -1,5 +1,7 @@ mod model; +use axum::http::header; +use axum::http::Method; use axum::{ response::{sse::Event, IntoResponse, Sse}, Json, @@ -7,7 +9,6 @@ use axum::{ use eventsource_stream::{EventStream, Eventsource}; use futures::StreamExt; use futures_core::Stream; -use http::header; use reqwest::StatusCode; use serde_json::Value; use std::{convert::Infallible, str::FromStr}; @@ -38,7 +39,7 @@ use crate::URL_CHATGPT_API; /// Check if the request is supported pub(super) fn support(req: &RequestExt) -> bool { - if req.uri.path().eq("/v1/chat/completions") && req.method.eq(&http::Method::POST) { + if req.uri.path().eq("/v1/chat/completions") && req.method.eq(&Method::POST) { if let Some(ref b) = req.bearer_auth() { return !b.starts_with("sk-"); } diff --git a/openai/src/serve/route/arkose.rs b/openai/src/serve/route/arkose.rs index 545773817..08e3d1baf 100644 --- a/openai/src/serve/route/arkose.rs +++ b/openai/src/serve/route/arkose.rs @@ -8,6 +8,7 @@ use crate::with_context; use axum::body::Body; use axum::http::header; use axum::http::method::Method; +use axum::http::response::Builder; use axum::http::Response; use axum::http::StatusCode; use axum::response::IntoResponse; @@ -17,7 +18,6 @@ use axum::{ Form, Router, }; use bytes::Bytes; -use http::response::Builder; use std::collections::HashMap; pub(super) fn config(router: Router, args: &ContextArgs) -> Router { diff --git a/openai/src/serve/route/files.rs b/openai/src/serve/route/files.rs index fd177055e..57ae8b5cf 100644 --- a/openai/src/serve/route/files.rs +++ b/openai/src/serve/route/files.rs @@ -1,5 +1,5 @@ +use axum::http::header; use axum::{response::IntoResponse, routing::any, Router}; -use http::header; use crate::{ context::ContextArgs, diff --git a/openai/src/serve/route/har/mod.rs b/openai/src/serve/route/har/mod.rs index beff1ce95..7416018a4 100644 --- a/openai/src/serve/route/har/mod.rs +++ b/openai/src/serve/route/har/mod.rs @@ -305,8 +305,8 @@ fn check_file_extension(file: &PathBuf) -> Result<(), Html> { } use axum::headers::{Header, HeaderName, HeaderValue}; +use axum::http::header; use axum_extra::extract::CookieJar; -use http::header; struct PlatformType(arkose::Type); diff --git a/openai/src/serve/route/ui/extract.rs b/openai/src/serve/route/ui/extract.rs index f4823494b..755a61119 100644 --- a/openai/src/serve/route/ui/extract.rs +++ b/openai/src/serve/route/ui/extract.rs @@ -1,9 +1,9 @@ use std::str::FromStr; +use axum::http::{HeaderMap, Request}; use axum::{async_trait, extract::FromRequest}; use axum_extra::extract::CookieJar; use base64::Engine; -use http::{HeaderMap, Request}; use serde::{Deserialize, Serialize}; use crate::{ diff --git a/openai/src/serve/route/ui/mod.rs b/openai/src/serve/route/ui/mod.rs index f9fd59d01..318a6ba06 100644 --- a/openai/src/serve/route/ui/mod.rs +++ b/openai/src/serve/route/ui/mod.rs @@ -9,6 +9,7 @@ use axum::extract::Query; use axum::headers::authorization::Bearer; use axum::headers::Authorization; use axum::http::header; +use axum::http::response::Builder; use axum::http::HeaderMap; use axum::http::Response; use axum::http::StatusCode; @@ -24,7 +25,6 @@ use axum_csrf::CsrfToken; use axum_csrf::Key; use axum_extra::extract::cookie; use axum_extra::extract::CookieJar; -use http::response::Builder; use serde_json::{json, Value}; use std::collections::HashMap; use std::net::SocketAddr; diff --git a/src/main.rs b/src/main.rs index 1b895229b..33db44637 100644 --- a/src/main.rs +++ b/src/main.rs @@ -31,7 +31,7 @@ fn main() -> anyhow::Result<()> { #[cfg(target_family = "unix")] args::ServeSubcommand::Log => handle::serve_log()?, args::ServeSubcommand::Genca => { - let _ = openai::serve::preauth::cagen::gen_ca(); + let _ = mitm::cagen::gen_ca(); } args::ServeSubcommand::GT { out } => handle::generate_template(out)?, args::ServeSubcommand::Update => update::update()?,