Plumb proxy into fetch and scrape (#88)

crates/obscura-cli/src/main.rs

@@ -19,7 +19,7 @@ struct Args {
     #[arg(short, long, default_value_t = 9222)]
     port: u16,
 
-    #[arg(long)]
+    #[arg(long, global = true)]
     proxy: Option<String>,
 
     #[arg(long)]
@@ -136,7 +136,14 @@ fn select_log_filter(verbose: bool, quiet: bool) -> &'static str {
 }
 
 fn is_quiet_command(cmd: &Option<Command>) -> bool {
-    matches!(cmd, Some(Command::Fetch { quiet: true, .. }))
+    matches!(
+        cmd,
+        Some(Command::Fetch { quiet: true, .. }) | Some(Command::Scrape { quiet: true, .. })
+    )
+}
+
+fn merge_proxy(global_proxy: Option<String>, command_proxy: Option<String>) -> Option<String> {
+    command_proxy.or(global_proxy)
 }
 
 #[tokio::main(flavor = "current_thread")]
@@ -153,8 +160,11 @@ async fn main() -> anyhow::Result<()> {
         .with_writer(std::io::stderr)
         .init();
 
+    let global_proxy = args.proxy.clone();
+
     match args.command {
         Some(Command::Serve { port, proxy, user_agent, stealth, workers }) => {
+            let proxy = merge_proxy(global_proxy.clone(), proxy);
             print_banner(port);
             if let Some(ref proxy) = proxy {
                 tracing::info!("Using proxy: {}", proxy);
@@ -179,10 +189,10 @@ async fn main() -> anyhow::Result<()> {
             }
         }
         Some(Command::Fetch { url, dump, selector, wait, timeout, wait_until, user_agent, stealth, eval, output, quiet }) => {
-            run_fetch(&url, dump, selector, wait, timeout, &wait_until, user_agent, stealth, eval, output, quiet).await?;
+            run_fetch(&url, dump, selector, wait, timeout, &wait_until, user_agent, stealth, eval, output, quiet, global_proxy).await?;
         }
         Some(Command::Scrape { urls, eval, concurrency, format, timeout, quiet }) => {
-            run_parallel_scrape(urls, eval, concurrency.get(), &format, timeout, quiet).await?;
+            run_parallel_scrape(urls, eval, concurrency.get(), &format, timeout, quiet, global_proxy).await?;
         }
         None => {
             print_banner(args.port);
@@ -335,8 +345,9 @@ async fn run_fetch(
     eval: Option<String>,
     output: Option<std::path::PathBuf>,
     quiet: bool,
+    proxy: Option<String>,
 ) -> anyhow::Result<()> {
-    let context = Arc::new(BrowserContext::with_options("fetch".to_string(), None, stealth));
+    let context = Arc::new(BrowserContext::with_options("fetch".to_string(), proxy, stealth));
     let mut page = Page::new("fetch-page".to_string(), context);
 
     if let Some(ref ua) = user_agent {
@@ -510,6 +521,7 @@ async fn run_parallel_scrape(
     format: &str,
     timeout_secs: u64,
     quiet: bool,
+    proxy: Option<String>,
 ) -> anyhow::Result<()> {
     let total = urls.len();
     let start = Instant::now();
@@ -550,6 +562,7 @@ async fn run_parallel_scrape(
         let sem = semaphore.clone();
         let eval = eval.clone();
         let worker_path = worker_path.clone();
+        let proxy = proxy.clone();
 
         let handle = tokio::spawn(async move {
             let _permit = sem.acquire().await.unwrap();
@@ -559,6 +572,7 @@ async fn run_parallel_scrape(
                 .stdin(std::process::Stdio::piped())
                 .stdout(std::process::Stdio::piped())
                 .stderr(std::process::Stdio::null())
+                .env("OBSCURA_PROXY", proxy.as_deref().unwrap_or(""))
                 .spawn()
             {
                 Ok(c) => c,
@@ -771,7 +785,8 @@ fn dump_links(page: &Page) -> String {
 #[cfg(test)]
 mod tests {
     use super::{
-        extract_readable_text, is_quiet_command, select_log_filter, write_or_print, Args, Command,
+        extract_readable_text, is_quiet_command, merge_proxy, select_log_filter, write_or_print,
+        Args, Command,
     };
     use clap::Parser;
     use obscura_dom::parse_html;
@@ -922,4 +937,21 @@ mod tests {
         assert!(!text.contains("console.log"));
         assert!(!text.contains("color: red"));
     }
+
+    #[test]
+    fn command_proxy_overrides_global_proxy() {
+        let proxy = merge_proxy(
+            Some("http://global.example:8080".to_string()),
+            Some("socks5://127.0.0.1:1080".to_string()),
+        );
+
+        assert_eq!(proxy.as_deref(), Some("socks5://127.0.0.1:1080"));
+    }
+
+    #[test]
+    fn global_proxy_is_used_when_command_proxy_is_absent() {
+        let proxy = merge_proxy(Some("http://global.example:8080".to_string()), None);
+
+        assert_eq!(proxy.as_deref(), Some("http://global.example:8080"));
+    }
 }