feat: add site config at cookies section

wesleyel · wesleyel · commit a209efae289a · 2024-07-04T03:17:23.000Z
diff --git a/README.md b/README.md
@@ -7,8 +7,6 @@
 [![telegram](https://img.shields.io/badge/telegram-blue?logo=telegram)](https://t.me/+U_5si6PhWykxZTI1)
 [![LICENSE](https://img.shields.io/crates/l/leetcode-cli.svg)](https://choosealicense.com/licenses/mit/)
 
-Modify for **leetcode.cn**
-
 ## Installing
 
 ```sh
@@ -45,7 +43,7 @@ If no argument is provided, the shell is inferred from the `SHELL` environment v
 
 ## Usage
 
-**Make sure you have logged in to `leetcode.cn` with `Firefox`**. See [Cookies](#cookies) for why you need to do this first.
+**Make sure you have logged in to `leetcode.com` with `Firefox`**. See [Cookies](#cookies) for why you need to do this first.
 
 ```sh
 leetcode 0.4.0
@@ -309,7 +307,7 @@ Open Firefox, press F12, and click `Storage` tab.
 
 #### Step 2
 
-Expand `Cookies` tab on the left and select https://leetcode.cn.
+Expand `Cookies` tab on the left and select https://leetcode.com.
 
 #### Step 2
 
diff --git a/src/cache/parser.rs b/src/cache/parser.rs
@@ -10,15 +10,17 @@ pub fn problem(problems: &mut Vec<Problem>, v: Value) -> Option<()> {
         let total_acs = stat.get("total_acs")?.as_f64()? as f32;
         let total_submitted = stat.get("total_submitted")?.as_f64()? as f32;
 
+        let fid_obj = stat.get("frontend_question_id")?;
+        let fid = match fid_obj.as_i64() {
+            // Handle on leetcode-com
+            Some(s) => s as i32,
+            // Handle on leetcode-cn
+            None => fid_obj.as_str()?.split(" ").last()?.parse::<i32>().ok()?,
+        };
+
         problems.push(Problem {
             category: v.get("category_slug")?.as_str()?.to_string(),
-            fid: stat
-                .get("frontend_question_id")?
-                .as_str()?
-                .split(" ")
-                .last()?
-                .parse::<i32>()
-                .ok()?,
+            fid: fid,
             id: stat.get("question_id")?.as_i64()? as i32,
             level: p.get("difficulty")?.as_object()?.get("level")?.as_i64()? as i32,
             locked: p.get("paid_only")?.as_bool()?,
@@ -95,18 +97,20 @@ pub fn tags(v: Value) -> Option<Vec<String>> {
 /// daily parser
 pub fn daily(v: Value) -> Option<i32> {
     trace!("Parse daily...");
-    v.as_object()?
-        .get("data")?
-        .as_object()?
-        .get("todayRecord")?
-        .as_array()?[0]
-        .as_object()?
-        .get("question")?
-        .as_object()?
-        .get("questionFrontendId")?
-        .as_str()?
-        .parse()
-        .ok()
+    let v_obj = v.as_object()?.get("data")?.as_object()?;
+    match v_obj.get("dailyQuestionRecord") {
+        // Handle on leetcode-com
+        Some(v) => v,
+        // Handle on leetcode-cn
+        None => v_obj.get("todayRecord")?.as_array()?.get(0)?,
+    }
+    .as_object()?
+    .get("question")?
+    .as_object()?
+    .get("questionFrontendId")?
+    .as_str()?
+    .parse()
+    .ok()
 }
 
 /// user parser
diff --git a/src/config/cookies.rs b/src/config/cookies.rs
@@ -1,11 +1,52 @@
 //! Cookies in config
+use std::str::FromStr;
+
 use serde::{Deserialize, Serialize};
 
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub enum LeetcodeSite {
+    #[serde(rename = "leetcode.com")]
+    LeetcodeCom,
+    #[serde(rename = "leetcode.cn")]
+    LeetcodeCn,
+}
+
+impl FromStr for LeetcodeSite {
+    type Err = String;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "leetcode.com" => Ok(LeetcodeSite::LeetcodeCom),
+            "leetcode.cn" => Ok(LeetcodeSite::LeetcodeCn),
+            _ => Err("Invalid site key".to_string()),
+        }
+    }
+}
+
+impl ToString for LeetcodeSite {
+    fn to_string(&self) -> String {
+        match self {
+            LeetcodeSite::LeetcodeCom => "leetcode.com".to_string(),
+            LeetcodeSite::LeetcodeCn => "leetcode.cn".to_string(),
+        }
+    }
+}
+
 /// Cookies settings
-#[derive(Default, Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct Cookies {
     pub csrf: String,
     pub session: String,
+    pub site: LeetcodeSite,
+}
+
+impl Default for Cookies {
+    fn default() -> Self {
+        Self {
+            csrf: "".to_string(),
+            session: "".to_string(),
+            site: LeetcodeSite::LeetcodeCom,
+        }
+    }
 }
 
 impl std::string::ToString for Cookies {
diff --git a/src/config/mod.rs b/src/config/mod.rs
@@ -17,6 +17,8 @@ mod cookies;
 mod storage;
 mod sys;
 
+pub use cookies::LeetcodeSite;
+
 /// Sync with `~/.leetcode/leetcode.toml`
 #[derive(Clone, Debug, Default, Deserialize, Serialize)]
 pub struct Config {
@@ -44,7 +46,14 @@ impl Config {
 
         let s = fs::read_to_string(&conf)?;
         match toml::from_str::<Config>(&s) {
-            Ok(config) => Ok(config),
+            Ok(config) => match config.cookies.site {
+                cookies::LeetcodeSite::LeetcodeCom => Ok(config),
+                cookies::LeetcodeSite::LeetcodeCn => {
+                    let mut config = config;
+                    config.sys.urls = sys::Urls::new_with_leetcode_cn();
+                    Ok(config)
+                }
+            },
             Err(e) => {
                 let tmp = Self::root()?.join("leetcode.tmp.toml");
                 Self::write_default(tmp)?;
diff --git a/src/config/sys.rs b/src/config/sys.rs
@@ -32,6 +32,27 @@ pub struct Urls {
 
 impl Default for Urls {
     fn default() -> Self {
+        Self {
+            base: "https://leetcode.com".into(),
+            graphql: "https://leetcode.com/graphql".into(),
+            login: "https://leetcode.com/accounts/login/".into(),
+            problems: "https://leetcode.com/api/problems/$category/".into(),
+            problem: "https://leetcode.com/problems/$slug/description/".into(),
+            tag: "https://leetcode.com/tag/$slug/".into(),
+            test: "https://leetcode.com/problems/$slug/interpret_solution/".into(),
+            session: "https://leetcode.com/session/".into(),
+            submit: "https://leetcode.com/problems/$slug/submit/".into(),
+            submissions: "https://leetcode.com/submissions/detail/$id/".into(),
+            submission: "https://leetcode.com/submissions/detail/$id/".into(),
+            verify: "https://leetcode.com/submissions/detail/$id/check/".into(),
+            favorites: "https://leetcode.com/list/api/questions".into(),
+            favorite_delete: "https://leetcode.com/list/api/questions/$hash/$id".into(),
+        }
+    }
+}
+
+impl Urls {
+    pub fn new_with_leetcode_cn() -> Self {
         Self {
             base: "https://leetcode.cn".into(),
             graphql: "https://leetcode.cn/graphql".into(),
@@ -49,9 +70,7 @@ impl Default for Urls {
             favorite_delete: "https://leetcode.cn/list/api/questions/$hash/$id".into(),
         }
     }
-}
 
-impl Urls {
     /// problem url with specific `$slug`
     pub fn problem(&self, slug: &str) -> String {
         self.problem.replace("$slug", slug)
diff --git a/src/lib.rs b/src/lib.rs
@@ -12,7 +12,7 @@
 //!
 //! ## Usage
 //!
-//! **Please make sure you have logined in `leetcode.cn` with `chrome`**, more info plz checkout [this](#cookies)
+//! **Please make sure you have logined in `leetcode.com` with `chrome`**, more info plz checkout [this](#cookies)
 //!
 //! ```sh
 //! leetcode 0.3.10
@@ -142,15 +142,15 @@
 //! session = "..."
 //! ```
 //!
-//! For Example, if you're using chrome to login to leetcode.cn.
+//! For Example, if you're using chrome to login to leetcode.com.
 //!
 //!
 //! #### Step 1
 //!
 //! Open chrome and paste the link below to the `chrome linkbar`.
 //!
 //! ```sh
-//! chrome://settings/cookies/detail?site=leetcode.cn
+//! chrome://settings/cookies/detail?site=leetcode.com
 //! ```
 //!
 //! #### Step 2
diff --git a/src/plugins/chrome.rs b/src/plugins/chrome.rs
@@ -64,7 +64,7 @@ pub fn cookies() -> Result<Ident> {
     debug!("Chrome Cookies path is {:?}", &p);
     let mut conn = cache::conn(p.to_string_lossy().to_string());
     let res = cookies
-        .filter(host_key.like("%leetcode.cn"))
+        .filter(host_key.like(format!("#{}", ccfg.site.to_string())))
         .load::<Cookies>(&mut conn)
         .expect("Loading cookies from google chrome failed.");
 
diff --git a/src/plugins/leetcode.rs b/src/plugins/leetcode.rs
@@ -145,21 +145,42 @@ impl LeetCode {
         trace!("Requesting daily problem...");
         let url = &self.conf.sys.urls.graphql;
         let mut json: Json = HashMap::new();
-        json.insert("operationName", "questionOfToday".to_string());
-        json.insert(
-            "query",
-            vec![
-                "query questionOfToday {",
-                "  todayRecord {",
-                "    question {",
-                "      questionFrontendId",
-                "    }",
-                "  }",
-                "}",
-            ]
-            .join("\n"),
-        );
-    
+
+        match self.conf.cookies.site {
+            config::LeetcodeSite::LeetcodeCom => {
+                json.insert("operationName", "daily".to_string());
+                json.insert(
+                    "query",
+                    vec![
+                        "query daily {",
+                        "  activeDailyCodingChallengeQuestion {",
+                        "    question {",
+                        "      questionFrontendId",
+                        "    }",
+                        "  }",
+                        "}",
+                    ]
+                    .join("\n"),
+                );
+            }
+            config::LeetcodeSite::LeetcodeCn => {
+                json.insert("operationName", "questionOfToday".to_string());
+                json.insert(
+                    "query",
+                    vec![
+                        "query questionOfToday {",
+                        "  todayRecord {",
+                        "    question {",
+                        "      questionFrontendId",
+                        "    }",
+                        "  }",
+                        "}",
+                    ]
+                    .join("\n"),
+                );
+            }
+        }
+
         Req {
             default_headers: self.default_headers,
             refer: None,
diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs
@@ -3,8 +3,8 @@
 //! + chrome cookie parser
 //! + leetcode API
 //!
-//! ## login to `leetcode.cn`
-//! Leetcode-cli use chrome cookie directly, do not need to login, please make sure you have loggined in `leetcode.cn` before usnig `leetcode-cli`
+//! ## login to `leetcode.com`
+//! Leetcode-cli use chrome cookie directly, do not need to login, please make sure you have loggined in `leetcode.com` before usnig `leetcode-cli`
 //!
 
 // FIXME: Read cookies from local storage. (issue #122)
