From c078f22826c46a39bc8f444e609eb910851e95db Mon Sep 17 00:00:00 2001 From: xsigoking <152482559+xsigoking@users.noreply.github.com> Date: Sat, 20 Apr 2024 11:22:14 +0000 Subject: [PATCH] feat: calcuate proof token (#14) --- Cargo.lock | 74 ++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/main.rs | 128 ++++++++++++++++++++++++++++------------------------ 3 files changed, 145 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1947d5e..13b0704 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -128,6 +128,15 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bumpalo" version = "3.16.0" @@ -173,6 +182,7 @@ dependencies = [ "reqwest-eventsource", "serde", "serde_json", + "sha3", "tokio", "tokio-graceful", "tokio-stream", @@ -215,6 +225,35 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f" +[[package]] +name = "cpufeatures" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53fe5e26ff1b7aef8bca9c6080520cfb8d9333c7568e1829cef191a9723e5504" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "either" version = "1.10.0" @@ -356,6 +395,16 @@ dependencies = [ "windows", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.14" @@ -578,6 +627,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "keccak" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc2af9a1119c51f12a14607e783cb977bde58bc069ff0c3da1095e635d70654" +dependencies = [ + "cpufeatures", +] + [[package]] name = "lazy_static" version = "1.4.0" @@ -1102,6 +1160,16 @@ dependencies = [ "serde", ] +[[package]] +name = "sha3" +version = "0.10.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60" +dependencies = [ + "digest", + "keccak", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -1406,6 +1474,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" + [[package]] name = "unicase" version = "2.7.0" diff --git a/Cargo.toml b/Cargo.toml index 9fa5112..3091065 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ rand = "0.8.5" reqwest-eventsource = "0.6.0" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0.68", features = ["preserve_order"] } +sha3 = "0.10.8" tokio = { version = "1.34.0", features = ["rt", "time", "macros", "rt-multi-thread"] } tokio-graceful = "0.1.6" tokio-stream = { version = "0.1.15", default-features = false, features = ["sync"] } diff --git a/src/main.rs b/src/main.rs index dfda5b2..9d8d3bb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,7 +4,7 @@ extern crate log; use anyhow::{anyhow, bail, Result}; use base64::{engine::general_purpose::STANDARD, Engine}; use bytes::Bytes; -use chrono::{DateTime, Datelike, Timelike, Utc}; +use chrono::Utc; use futures_util::StreamExt; use http::{HeaderMap, HeaderValue, Response, StatusCode}; use http_body_util::{combinators::BoxBody, BodyExt, Full, StreamBody}; @@ -17,6 +17,7 @@ use rand::{seq::SliceRandom, thread_rng, Rng}; use reqwest::{Client, ClientBuilder, Method, Proxy}; use reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt}; use serde_json::{json, Value}; +use sha3::{Digest, Sha3_512}; use std::{convert::Infallible, env, sync::Arc, time::Duration}; use tokio::{ net::TcpListener, @@ -41,10 +42,6 @@ lazy_static::lazy_static! { let mut rng = rand::thread_rng(); rng.gen_range(2000..8000) }; - static ref PROOF_V2: u32 = { - let mut rng = rand::thread_rng(); - rng.gen_range(8..24) - }; } #[tokio::main] @@ -197,7 +194,7 @@ impl Server { } async fn chat_completion(&self, req: hyper::Request) -> Result { - let (oai_device_id, token) = self + let requirements = self .chat_requirements() .await .map_err(|err| anyhow!("Failed to meet chat requirements, {err}"))?; @@ -265,16 +262,19 @@ impl Server { "websocket_request_id": random_id(), }); - let proof_token = openai_sentinel_proof_token(); - debug!("headers: oai_device_id {oai_device_id}; openai-sentinel-chat-requirements-token {token}; openai-sentinel-proof-token {proof_token}"); + let proof_token = calculate_proof_token(&requirements.seed, &requirements.difficulty); + debug!("headers: oai_device_id {}; openai-sentinel-chat-requirements-token {}; openai-sentinel-proof-token {proof_token}", requirements.oai_device_id, requirements.token); debug!("req body: {req_body}"); let mut es = self .client .post(CONVERSATION_URL) .headers(common_headers()) - .header("oai-device-id", oai_device_id) - .header("openai-sentinel-chat-requirements-token", token) + .header("oai-device-id", requirements.oai_device_id) + .header( + "openai-sentinel-chat-requirements-token", + requirements.token, + ) .header("openai-sentinel-proof-token", proof_token) .json(&req_body) .eventsource()?; @@ -424,7 +424,7 @@ impl Server { Ok(res) } - async fn chat_requirements(&self) -> Result<(String, String)> { + async fn chat_requirements(&self) -> Result { let oai_device_id = random_id(); let res = self .client @@ -435,11 +435,27 @@ impl Server { .send() .await?; let data: Value = res.json().await?; - let token = match data["token"].as_str() { - Some(v) => v.to_string(), - None => bail!("Invalid data, {data}"), - }; - Ok((oai_device_id, token)) + if let (Some(token), Some((seed, difficulty))) = ( + data["token"].as_str(), + data["proofofwork"].as_object().and_then(|v| { + if let (Some(seed), Some(difficulty)) = + (v["seed"].as_str(), v["difficulty"].as_str()) + { + Some((seed, difficulty)) + } else { + None + } + }), + ) { + Ok(Requirements { + oai_device_id, + token: token.to_string(), + seed: seed.to_string(), + difficulty: difficulty.to_string(), + }) + } else { + bail!("Invalid data, {data}"); + } } } @@ -450,6 +466,14 @@ enum ResEvent { Done, } +#[derive(Debug)] +struct Requirements { + oai_device_id: String, + token: String, + seed: String, + difficulty: String, +} + async fn send_first_event(tx: Sender, data: Option, check: &mut bool) { if *check { let _ = tx.send(ResEvent::First(data)).await; @@ -490,6 +514,7 @@ fn common_headers() -> HeaderMap { HeaderValue::from_static("https://chat.openai.com"), ); headers.insert("pragma", HeaderValue::from_static("no-cache")); + headers.insert("priority", HeaderValue::from_static("u=1, i")); headers.insert( "referer", HeaderValue::from_static("https://chat.openai.com/"), @@ -609,49 +634,36 @@ fn random_id() -> String { Uuid::new_v4().to_string() } -fn openai_sentinel_proof_token() -> String { - let datetime = format_date_time(&Utc::now()); - let value = format!( - r#"[{},"{datetime}",4294705152,{},"{USER_AGENT}"]"#, - *PROOF_V1, *PROOF_V2 - ); - let value = STANDARD.encode(value); - format!("gAAAAAB{value}") -} +fn calculate_proof_token(seed: &str, diff: &str) -> String { + let now = Utc::now(); + let datetime = now.format("%a %b %d %Y %H:%M:%S GMT%z (Coordinated Universal Time)"); + + let diff_len = diff.len() / 2; + let mut hasher = Sha3_512::new(); + + for i in 0..100000 { + let value = format!( + r#"[{},"{datetime}",4294705152,{},"{USER_AGENT}"]"#, + *PROOF_V1, i + ); + let base = STANDARD.encode(value); + hasher.update(format!("{}{}", seed, base).as_bytes()); + let hash = hasher.finalize_reset(); + let hash_hex = hex_encode(&hash[..diff_len]); + + if hash_hex.as_str() <= diff { + return format!("gAAAAAB{}", base); + } + } -fn format_date_time(dt: &DateTime) -> String { - let weekday = match dt.weekday() { - chrono::Weekday::Sun => "Sun", - chrono::Weekday::Mon => "Mon", - chrono::Weekday::Tue => "Tue", - chrono::Weekday::Wed => "Wed", - chrono::Weekday::Thu => "Thu", - chrono::Weekday::Fri => "Fri", - chrono::Weekday::Sat => "Sat", - }; - let month = match dt.month() { - 1 => "Jan", - 2 => "Feb", - 3 => "Mar", - 4 => "Apr", - 5 => "May", - 6 => "Jun", - 7 => "Jul", - 8 => "Aug", - 9 => "Sep", - 10 => "Oct", - 11 => "Nov", - 12 => "Dec", - _ => unreachable!(), - }; format!( - "{} {} {:02} {} {:02}:{:02}:{:02} GMT+0000 (Coordinated Universal Time)", - weekday, - month, - dt.day(), - dt.year(), - dt.hour(), - dt.minute(), - dt.second(), + "gAAAAABwQ8Lk5FbGpA2NcR9dShT6gYjU7VxZ4D{}", + STANDARD.encode(format!("\"{}\"", seed)) ) } + +fn hex_encode(bytes: &[u8]) -> String { + bytes + .iter() + .fold(String::new(), |acc, b| acc + &format!("{:02x}", b)) +}