From ab60ba29e9c7729e7356f287c8eb5b28355bb64f Mon Sep 17 00:00:00 2001
From: Nick Carton <nick.carton@gmail.com>
Date: Sat, 30 Nov 2024 18:05:16 +0100
Subject: [PATCH] Allow specifying proxy, residential or mobile connections

---
 src/main.rs         | 10 ++++++---
 src/options/opts.rs | 53 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 31abc61..a4eb835 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -122,10 +122,14 @@ async fn main() -> eyre::Result<()> {
                             hostnames: vec![endpoint.to_string()],
                             isp_regex: None,
                             city: None,
-                            mobile: PerformIcmpBodyMobile::from_str("ALLOWED")?,
+                            mobile: PerformIcmpBodyMobile::from_str(
+                                &APP_CONFIG.mobile.to_string(),
+                            )?,
                             node_id: None,
-                            proxy: PerformIcmpBodyProxy::from_str("ALLOWED")?,
-                            residential: PerformIcmpBodyResidential::from_str("ALLOWED")?,
+                            proxy: PerformIcmpBodyProxy::from_str(&APP_CONFIG.proxy.to_string())?,
+                            residential: PerformIcmpBodyResidential::from_str(
+                                &APP_CONFIG.residential.to_string(),
+                            )?,
                         })
                         .await
                         .context("Failed to send job");
diff --git a/src/options/opts.rs b/src/options/opts.rs
index 6c8c767..3c8ba6c 100644
--- a/src/options/opts.rs
+++ b/src/options/opts.rs
@@ -1,7 +1,36 @@
+use std::fmt::Display;
+
 use bpaf::{long, OptionParser, Parser};
 use color_eyre::eyre;
 use keshvar::Continent;
 
+#[derive(Debug, Clone)]
+pub enum NetworkPolicy {
+    Allowed,
+    Denied,
+    Required,
+}
+
+impl From<Option<bool>> for NetworkPolicy {
+    fn from(_value: Option<bool>) -> Self {
+        match _value {
+            None => NetworkPolicy::Allowed,
+            Some(true) => NetworkPolicy::Required,
+            Some(false) => NetworkPolicy::Denied,
+        }
+    }
+}
+
+impl Display for NetworkPolicy {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            NetworkPolicy::Allowed => write!(f, "ALLOWED"),
+            NetworkPolicy::Denied => write!(f, "DENIED"),
+            NetworkPolicy::Required => write!(f, "REQUIRED"),
+        }
+    }
+}
+
 /// This doc string acts as a help message when the user runs '--help'
 /// as do all doc strings on fields
 #[derive(Debug, Clone)]
@@ -12,6 +41,9 @@ pub struct Opts {
     pub attempts: usize,
     pub api_key: String,
     pub concurrency: usize,
+    pub residential: NetworkPolicy,
+    pub mobile: NetworkPolicy,
+    pub proxy: NetworkPolicy,
 }
 
 impl Opts {
@@ -53,12 +85,33 @@ impl Opts {
             .argument::<usize>("concurrency")
             .fallback(100);
 
+        let residential = bpaf::long("residential")
+            .help("Control residential network usage. --residential=true to require, --residential=false to deny, omit to allow.")
+            .argument::<bool>("residential")
+            .optional()
+            .map(NetworkPolicy::from);
+
+        let mobile = bpaf::long("mobile")
+            .help("Control mobile network usage. --mobile=true to require, --mobile=false to deny, omit to allow.")
+            .argument::<bool>("mobile")
+            .optional()
+            .map(NetworkPolicy::from);
+
+        let proxy = bpaf::long("proxy")
+            .help("Control proxy network usage. --proxy=true to require, --proxy=false to deny, omit to allow.")
+            .argument::<bool>("proxy")
+            .optional()
+            .map(NetworkPolicy::from);
+
         bpaf::construct!(Opts {
             regions,
             count,
             attempts,
             concurrency,
             api_key,
+            residential,
+            mobile,
+            proxy,
             endpoint,
         })
         .to_options()