feat: update connect.rs

connect-rs/src/main.rs

@@ -15,7 +15,8 @@ use crypto::digest::Digest;
 use crypto::sha2::Sha256;
 use keyring::Entry;
 use rust_util::{
-    XResult, failure, information, opt_result, simple_error, success, util_cmd, util_env, util_file,
+    failure, iff, information, opt_result, simple_error, success, util_cmd, util_env, util_file,
+    XResult,
 };
 use serde_json::Value;
 use std::fs;
@@ -24,6 +25,12 @@ use std::process::Command;
 const PAC_FILE: &str = "__proxy__.pac";
 const ENV_CONNECT_PAC_FILE_PATH: &str = "CONNECT_PAC_FILE_PATH";
 
+const KEYRING_SERVICE: &str = "hatter.ink";
+const KEYRING_USER: &str = "pac_auth_token";
+
+const NET_WORK_SETUP: &str = "networksetup";
+const WIFI: &str = "Wi-Fi";
+
 #[derive(Debug, Parser)]
 #[command(name = "connect-rs", bin_name = "connect.rs")]
 #[command(about = "Connect to the world", long_about = None)]
@@ -43,6 +50,9 @@ enum Commands {
     /// Update
     #[command(short_flag = 'u')]
     Update,
+    /// Add domain
+    #[command(short_flag = 'a')]
+    AddDomain(CmdAddDomain),
 }
 
 #[derive(Debug, Args)]
@@ -52,8 +62,20 @@ pub struct CmdConnect {
     pub update: bool,
 }
 
-const NET_WORK_SETUP: &str = "networksetup";
-const WIFI: &str = "Wi-Fi";
+#[derive(Debug, Args)]
+pub struct CmdAddDomain {
+    /// Domains (split by ',')
+    #[arg(long, short = 'd')]
+    pub domains: String,
+
+    /// Not GFW
+    #[arg(long)]
+    pub not_gfw: bool,
+
+    /// Reconnect
+    #[arg(long, short = 'r')]
+    pub reconnect: bool,
+}
 
 fn main() {
     let args = ConnectArgs::parse();
@@ -75,6 +97,16 @@ fn main() {
         Commands::Update => {
             update_pac_file_ignore_error();
         }
+        Commands::AddDomain(cmd) => match add_pac_digest(&cmd.domains, !cmd.not_gfw) {
+            Ok(_) => {
+                success!("Add domain successfully");
+                if cmd.reconnect {
+                    turn_proxy_off();
+                    turn_proxy_on();
+                }
+            }
+            Err(e) => failure!("Add domain failed: {}", e),
+        },
     }
 }
 
@@ -117,6 +149,14 @@ fn update_pac_file() -> XResult<()> {
     Ok(())
 }
 
+fn turn_proxy_on() {
+    run_command(
+        NET_WORK_SETUP,
+        &["-setautoproxyurl", WIFI, "http://localhost:8888/pac"],
+        "setautoproxyurl proxy",
+    );
+}
+
 fn turn_proxy_off() {
     run_command(
         NET_WORK_SETUP,
@@ -214,9 +254,29 @@ fn get_pac_digest() -> XResult<String> {
     }
 }
 
+fn add_pac_digest(domains: &str, gfw: bool) -> XResult<()> {
+    let auth_token = get_pac_auth_token()?;
+    let url = format!(
+        "https://hatter.ink/pac/add_pac_domain.json?digest=1&__auth_token={}&type={}&domains={}",
+        auth_token,
+        iff!(gfw, "gfw", "no_gfw"),
+        domains
+    );
+
+    let pac_digest_response =
+        opt_result!(reqwest::blocking::get(&url), "Add domain to pac failed: {}");
+    if pac_digest_response.status().as_u16() != 200 {
+        return simple_error!(
+            "Get pac digest HTTP status error: {}",
+            pac_digest_response.status().as_u16()
+        );
+    }
+    Ok(())
+}
+
 fn get_pac_auth_token() -> XResult<String> {
     let entry = opt_result!(
-        Entry::new("hatter.ink", "get_pac_auth_token"),
+        Entry::new(KEYRING_SERVICE, KEYRING_USER),
         "New entry failed: {}"
     );
     Ok(opt_result!(
@@ -225,5 +285,5 @@ fn get_pac_auth_token() -> XResult<String> {
     ))
 }
 
-// @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20251019T230159+08:00.MEQCIByk8UrJGAfSk3tNh4cr
-// +0M5B6P+72yivwmFvGnWOzv4AiBAjVanUpEEU4hUxpP8I3B/xMEdm26eDe25/7VAaxEmKw==
+// @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20251021T133159+08:00.MEQCIFEij7e8qbFLQ2oD7+uw
+// mvdUy6nMUPIb/cxG6UDKm2ccAiBHgFSSqZ9DXdwPnn/yZH6vJVyqvricLpTuRApNT9g+yA==