From 3b9c9ab7c9631d4915b780371d22d2a9227b1dca Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Mon, 6 Feb 2023 04:58:17 +0000
Subject: [PATCH] Add new option: user-agent

---
 src/fanbox_api.rs    |  3 ++-
 src/opthelper.rs     | 12 ++++++++++++
 src/opts.rs          |  4 ++++
 src/pixiv_web.rs     |  2 +-
 src/settings_list.rs |  1 +
 5 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/fanbox_api.rs b/src/fanbox_api.rs
index 157dac5..7ca4e83 100644
--- a/src/fanbox_api.rs
+++ b/src/fanbox_api.rs
@@ -84,7 +84,8 @@ impl FanboxClientInternal {
             }
             None => {}
         }
-        self.client.set_header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36");
+        let helper = get_helper();
+        self.client.set_header("User-Agent", &helper.user_agent());
         self.client.set_header("referer", "https://www.fanbox.cc/");
         self.client.set_header("origin", "https://www.fanbox.cc");
         self.inited.qstore(true);
diff --git a/src/opthelper.rs b/src/opthelper.rs
index 9502b33..bb00aef 100644
--- a/src/opthelper.rs
+++ b/src/opthelper.rs
@@ -376,6 +376,18 @@ impl OptHelper {
         false
     }
 
+    pub fn user_agent(&self) -> String {
+        match self.opt.get_ref().user_agent.as_ref() {
+            Some(ua) => return ua.to_owned(),
+            None => {}
+        }
+        match self.settings.get_ref().get("user-agent") {
+            Some(ua) => return ua.as_str().unwrap().to_owned(),
+            None => {}
+        }
+        String::from("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36")
+    }
+
     #[cfg(feature = "exif")]
     /// Return whether to add/update exif information to image files even
     /// when overwrite are disabled.
diff --git a/src/opts.rs b/src/opts.rs
index 6b2a317..7abff50 100644
--- a/src/opts.rs
+++ b/src/opts.rs
@@ -110,6 +110,7 @@ pub struct CommandOpts {
     pub x264_profile: Option<X264Profile>,
     /// The base directory to save downloaded files.
     pub download_base: Option<String>,
+    pub user_agent: Option<String>,
     /// Urls want to download
     pub urls: Option<Vec<String>>,
 }
@@ -148,6 +149,7 @@ impl CommandOpts {
             #[cfg(feature = "ugoira")]
             x264_profile: None,
             download_base: None,
+            user_agent: None,
             urls: None,
         }
     }
@@ -530,6 +532,7 @@ pub fn parse_cmd() -> Option<CommandOpts> {
         gettext("The base directory to download files."),
         "DIR",
     );
+    opts.optopt("", "user-agent", gettext("The User-Agent header."), "UA");
     let result = match opts.parse(&argv[1..]) {
         Ok(m) => m,
         Err(err) => {
@@ -826,6 +829,7 @@ pub fn parse_cmd() -> Option<CommandOpts> {
         }
     }
     re.as_mut().unwrap().download_base = result.opt_str("download-base");
+    re.as_mut().unwrap().user_agent = result.opt_str("user-agent");
     re
 }
 
diff --git a/src/pixiv_web.rs b/src/pixiv_web.rs
index 196f6d1..5d7689f 100644
--- a/src/pixiv_web.rs
+++ b/src/pixiv_web.rs
@@ -42,7 +42,7 @@ impl PixivWebClient {
                 return false;
             }
         }
-        self.client.set_header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36");
+        self.client.set_header("User-Agent", &helper.user_agent());
         let l = helper.language();
         if l.is_some() {
             self.client
diff --git a/src/settings_list.rs b/src/settings_list.rs
index fec5b4f..b5f88a9 100644
--- a/src/settings_list.rs
+++ b/src/settings_list.rs
@@ -54,6 +54,7 @@ pub fn get_settings_list() -> Vec<SettingDes> {
         #[cfg(feature = "db")]
         SettingDes::new("db", gettext("Database settings."), JsonValueType::Object, Some(check_db_config)).unwrap(),
         SettingDes::new("download-base", gettext("The base directory to save downloaded files."), JsonValueType::Str, None).unwrap(),
+        SettingDes::new("user-agent", gettext("The User-Agent header."), JsonValueType::Str, None).unwrap(),
     ]
 }