feat: initial commit — antigravity proxy with MITM, standalone LS, and snapshot tooling
This commit is contained in:
Nikketryhard
332
src/main.rs
Normal file
332
src/main.rs
Normal file
@@ -0,0 +1,332 @@
|
||||
//! Antigravity OpenAI Proxy — Rust edition v3 (stealth hardened).
|
||||
//!
|
||||
//! Single-binary replacement for server.py. BoringSSL TLS impersonation,
|
||||
//! byte-exact protobuf encoding, Chrome header fingerprinting, cascade
|
||||
//! session management, warmup + heartbeat lifecycle mimicry.
|
||||
|
||||
mod api;
|
||||
mod backend;
|
||||
mod constants;
|
||||
mod mitm;
|
||||
mod proto;
|
||||
mod quota;
|
||||
mod session;
|
||||
mod warmup;
|
||||
|
||||
use api::AppState;
|
||||
use backend::Backend;
|
||||
use clap::Parser;
|
||||
use session::SessionManager;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use mitm::store::MitmStore;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "antigravity-proxy", about = "Antigravity OpenAI Proxy (stealth)")]
|
||||
struct Cli {
|
||||
/// Port to listen on
|
||||
#[arg(long, default_value_t = 8741)]
|
||||
port: u16,
|
||||
|
||||
/// Enable info-level logging (-v)
|
||||
#[arg(short, long)]
|
||||
verbose: bool,
|
||||
|
||||
/// Enable debug-level logging (-d)
|
||||
#[arg(short, long)]
|
||||
debug: bool,
|
||||
|
||||
/// Disable the MITM proxy (no API interception)
|
||||
#[arg(long)]
|
||||
no_mitm: bool,
|
||||
|
||||
/// MITM proxy port (default: 8742, matches wrapper script)
|
||||
#[arg(long, default_value_t = 8742)]
|
||||
mitm_port: u16,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let cli = Cli::parse();
|
||||
|
||||
// Flag > env var > default (warn)
|
||||
let log_level = if cli.debug {
|
||||
"debug"
|
||||
} else if cli.verbose {
|
||||
"info"
|
||||
} else {
|
||||
// Fall back to RUST_LOG env, or warn-only
|
||||
""
|
||||
};
|
||||
|
||||
let filter = if log_level.is_empty() {
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "warn".into())
|
||||
} else {
|
||||
tracing_subscriber::EnvFilter::new(log_level)
|
||||
};
|
||||
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(filter)
|
||||
.init();
|
||||
|
||||
// ── Step 1: Bind main port FIRST (fail fast, before spawning anything) ────
|
||||
let addr = format!("127.0.0.1:{}", cli.port);
|
||||
let listener = match tokio::net::TcpListener::bind(&addr).await {
|
||||
Ok(l) => l,
|
||||
Err(e) => {
|
||||
eprintln!("Fatal: cannot bind to {addr}: {e}");
|
||||
eprintln!("Hint: kill $(lsof -ti:{}) 2>/dev/null", cli.port);
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
|
||||
// ── Step 2: Backend discovery ─────────────────────────────────────────────
|
||||
let backend = Arc::new(match Backend::new() {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
eprintln!("Fatal: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
});
|
||||
|
||||
let (pid, https_port, csrf, token) = backend.info().await;
|
||||
|
||||
// ── Step 3: MITM proxy (after port is secured) ────────────────────────────
|
||||
let mitm_store = MitmStore::new();
|
||||
let (mitm_port_actual, mitm_handle) = if !cli.no_mitm {
|
||||
let data_dir = dirs_data_dir();
|
||||
match mitm::ca::MitmCa::load_or_generate(&data_dir) {
|
||||
Ok(ca) => {
|
||||
let ca = Arc::new(ca);
|
||||
let ca_pem = ca.ca_pem_path.display().to_string();
|
||||
let config = mitm::proxy::MitmConfig {
|
||||
port: cli.mitm_port,
|
||||
modify_requests: false,
|
||||
};
|
||||
match mitm::proxy::run(ca, mitm_store.clone(), config).await {
|
||||
Ok((port, handle)) => {
|
||||
info!(port, ca = %ca_pem, "MITM proxy started");
|
||||
// Write actual port to file for wrapper script discovery
|
||||
let port_file = data_dir.join("mitm-port");
|
||||
if let Err(e) = std::fs::write(&port_file, port.to_string()) {
|
||||
warn!("Failed to write MITM port file: {e}");
|
||||
}
|
||||
(Some((port, ca_pem)), Some(handle))
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("MITM proxy failed to start: {e}");
|
||||
(None, None)
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("MITM CA generation failed: {e}");
|
||||
(None, None)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
info!("MITM proxy disabled (--no-mitm)");
|
||||
(None, None)
|
||||
};
|
||||
|
||||
// ── Step 4: Warmup + heartbeat ────────────────────────────────────────────
|
||||
warmup::warmup_sequence(&backend).await;
|
||||
let heartbeat_handle = warmup::start_heartbeat(Arc::clone(&backend));
|
||||
|
||||
// ── Step 4b: Quota monitor ────────────────────────────────────────────────
|
||||
let quota_store = quota::QuotaStore::new();
|
||||
quota_store.clone().start_polling(Arc::clone(&backend));
|
||||
info!("Quota monitor started (polling every 60s)");
|
||||
|
||||
let state = Arc::new(AppState {
|
||||
backend,
|
||||
sessions: SessionManager::new(),
|
||||
mitm_store,
|
||||
quota_store,
|
||||
});
|
||||
|
||||
// Periodic backend refresh — keeps LS connection details fresh
|
||||
let refresh_backend = Arc::clone(&state.backend);
|
||||
let refresh_handle = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
|
||||
if let Err(e) = refresh_backend.refresh().await {
|
||||
warn!("Periodic refresh failed: {e}");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// ── Step 5: Start serving ─────────────────────────────────────────────────
|
||||
let app = api::router(state.clone());
|
||||
|
||||
print_banner(cli.port, &pid, &https_port, &csrf, &token, &mitm_port_actual);
|
||||
info!("Listening on http://{addr}");
|
||||
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await
|
||||
.expect("server error");
|
||||
|
||||
// ── Cleanup: abort all background tasks ───────────────────────────────────
|
||||
heartbeat_handle.abort();
|
||||
refresh_handle.abort();
|
||||
if let Some(h) = mitm_handle {
|
||||
h.abort();
|
||||
}
|
||||
// Remove stale MITM port file
|
||||
let _ = std::fs::remove_file(dirs_data_dir().join("mitm-port"));
|
||||
info!("Server shutdown complete");
|
||||
}
|
||||
|
||||
/// Wait for SIGINT (Ctrl+C) or SIGTERM for graceful shutdown.
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
tokio::signal::ctrl_c()
|
||||
.await
|
||||
.expect("failed to install Ctrl+C handler");
|
||||
};
|
||||
|
||||
#[cfg(unix)]
|
||||
let terminate = async {
|
||||
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
||||
.expect("failed to install SIGTERM handler")
|
||||
.recv()
|
||||
.await;
|
||||
};
|
||||
|
||||
#[cfg(not(unix))]
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => info!("Received SIGINT, shutting down..."),
|
||||
_ = terminate => info!("Received SIGTERM, shutting down..."),
|
||||
}
|
||||
}
|
||||
|
||||
fn print_banner(port: u16, pid: &str, https_port: &str, csrf: &str, token: &str, mitm: &Option<(u16, String)>) {
|
||||
let chrome_major = &*constants::CHROME_MAJOR;
|
||||
let ver = crate::constants::antigravity_version();
|
||||
|
||||
println!();
|
||||
println!(" \x1b[1;35m>> antigravity-proxy\x1b[0m \x1b[2mv{ver}\x1b[0m");
|
||||
println!(" \x1b[2m────────────────────────────────────────────────\x1b[0m");
|
||||
println!();
|
||||
println!(" \x1b[1mcore\x1b[0m");
|
||||
println!(" \x1b[36m tls\x1b[0m BoringSSL (Chrome {chrome_major})");
|
||||
println!(" \x1b[36m listen\x1b[0m http://127.0.0.1:{port}");
|
||||
println!(" \x1b[36m ls pid\x1b[0m {pid}");
|
||||
println!(" \x1b[36m https\x1b[0m :{https_port}");
|
||||
println!(" \x1b[36m csrf\x1b[0m {csrf}");
|
||||
println!(" \x1b[36m oauth\x1b[0m {token}");
|
||||
println!();
|
||||
|
||||
// MITM section
|
||||
if let Some((mitm_port, ca_path)) = mitm {
|
||||
println!(" \x1b[1mmitm\x1b[0m");
|
||||
println!(" \x1b[36m proxy\x1b[0m 127.0.0.1:{mitm_port}");
|
||||
println!(" \x1b[36m ca cert\x1b[0m {ca_path}");
|
||||
|
||||
// Check if wrapper is installed
|
||||
let wrapper_installed = check_wrapper_installed();
|
||||
if wrapper_installed {
|
||||
println!(" \x1b[36m wrapper\x1b[0m \x1b[32minstalled\x1b[0m");
|
||||
} else {
|
||||
println!(" \x1b[36m wrapper\x1b[0m \x1b[33mnot installed\x1b[0m");
|
||||
}
|
||||
println!();
|
||||
} else {
|
||||
println!(" \x1b[1mmitm\x1b[0m \x1b[33mdisabled\x1b[0m");
|
||||
println!();
|
||||
}
|
||||
|
||||
// Routes
|
||||
println!(" \x1b[1mroutes\x1b[0m");
|
||||
println!(" \x1b[33m POST\x1b[0m /v1/responses");
|
||||
println!(" \x1b[33m POST\x1b[0m /v1/chat/completions");
|
||||
println!(" \x1b[32m GET \x1b[0m /v1/models");
|
||||
println!(" \x1b[32m GET \x1b[0m /v1/sessions");
|
||||
println!(" \x1b[31m DEL \x1b[0m /v1/sessions/:id");
|
||||
println!(" \x1b[33m POST\x1b[0m /v1/token");
|
||||
println!(" \x1b[32m GET \x1b[0m /v1/usage");
|
||||
println!(" \x1b[32m GET \x1b[0m /v1/quota");
|
||||
println!(" \x1b[32m GET \x1b[0m /health");
|
||||
println!();
|
||||
|
||||
// Status line
|
||||
let mitm_tag = if mitm.is_some() { "\x1b[32mmitm\x1b[0m" } else { "\x1b[31mmitm\x1b[0m" };
|
||||
println!(" \x1b[2mstealth:\x1b[0m \x1b[32mwarmup\x1b[0m \x1b[32mheartbeat\x1b[0m \x1b[32mjitter\x1b[0m {mitm_tag}");
|
||||
println!();
|
||||
|
||||
// Setup hints
|
||||
if let Some((mitm_port, ca_path)) = mitm {
|
||||
if !check_wrapper_installed() {
|
||||
println!(" \x1b[1;33m[!]\x1b[0m mitm wrapper not installed");
|
||||
println!(" \x1b[2mrun:\x1b[0m ./scripts/mitm-wrapper.sh install");
|
||||
println!(" \x1b[2mor:\x1b[0m HTTPS_PROXY=http://127.0.0.1:{mitm_port}");
|
||||
println!(" NODE_EXTRA_CA_CERTS={ca_path}");
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
if token == "NOT SET" {
|
||||
println!(" \x1b[1;33m[!]\x1b[0m no oauth token");
|
||||
println!(" export ANTIGRAVITY_OAUTH_TOKEN=ya29.xxx");
|
||||
println!(" curl -X POST http://127.0.0.1:{port}/v1/token -d '{{\"token\":\"ya29.xxx\"}}'");
|
||||
println!(" echo 'ya29.xxx' > ~/.config/antigravity-proxy-token");
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the MITM wrapper is installed by looking for the .real backup file
|
||||
/// next to the LS binary. Uses /proc to find the real LS path dynamically.
|
||||
fn check_wrapper_installed() -> bool {
|
||||
// Find the LS binary path from known PID or by scanning /proc
|
||||
if let Some(ls_path) = find_ls_binary_path() {
|
||||
let real_path = format!("{ls_path}.real");
|
||||
return std::path::Path::new(&real_path).exists();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Find the LS binary path by reading /proc/<pid>/exe for known language server processes.
|
||||
fn find_ls_binary_path() -> Option<String> {
|
||||
// Try all running processes, look for ones that look like the LS
|
||||
let proc = std::path::Path::new("/proc");
|
||||
if !proc.exists() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Ok(entries) = std::fs::read_dir(proc) {
|
||||
for entry in entries.flatten() {
|
||||
let name = entry.file_name();
|
||||
let name_str = name.to_string_lossy();
|
||||
// Only look at numeric dirs (PIDs)
|
||||
if !name_str.chars().all(|c| c.is_ascii_digit()) {
|
||||
continue;
|
||||
}
|
||||
let exe_link = entry.path().join("exe");
|
||||
if let Ok(target) = std::fs::read_link(&exe_link) {
|
||||
let target_str = target.to_string_lossy();
|
||||
// Strip " (deleted)" suffix from unlinked binaries
|
||||
let target_clean = target_str.trim_end_matches(" (deleted)");
|
||||
// Match any binary that looks like the Antigravity LS
|
||||
if target_clean.contains("language_server_linux")
|
||||
|| target_clean.contains("antigravity-language-server")
|
||||
{
|
||||
// Strip .real suffix — if the wrapper exec'd the backup, we want the base name
|
||||
let path = target_clean.trim_end_matches(".real");
|
||||
return Some(path.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Get the data directory for storing MITM CA cert/key.
|
||||
fn dirs_data_dir() -> std::path::PathBuf {
|
||||
let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
|
||||
std::path::PathBuf::from(home).join(".config").join("antigravity-proxy")
|
||||
}
|
||||
Reference in New Issue
Block a user