leetcode_cli/config/
cookies.rs

1//! Cookies in config
2use serde::{Deserialize, Serialize};
3use std::{
4    fmt::{self, Display},
5    str::FromStr,
6};
7
8#[derive(Clone, Debug, Deserialize, Serialize)]
9pub enum LeetcodeSite {
10    #[serde(rename = "leetcode.com")]
11    LeetcodeCom,
12    #[serde(rename = "leetcode.cn")]
13    LeetcodeCn,
14}
15
16impl FromStr for LeetcodeSite {
17    type Err = String;
18    fn from_str(s: &str) -> Result<Self, Self::Err> {
19        match s {
20            "leetcode.com" => Ok(LeetcodeSite::LeetcodeCom),
21            "leetcode.cn" => Ok(LeetcodeSite::LeetcodeCn),
22            _ => Err("Invalid site key".to_string()),
23        }
24    }
25}
26
27impl Display for LeetcodeSite {
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        let s = match self {
30            LeetcodeSite::LeetcodeCom => "leetcode.com",
31            LeetcodeSite::LeetcodeCn => "leetcode.cn",
32        };
33
34        write!(f, "{s}")
35    }
36}
37
38/// Cookies settings
39#[derive(Clone, Debug, Deserialize, Serialize)]
40pub struct Cookies {
41    pub csrf: String,
42    pub session: String,
43    pub site: LeetcodeSite,
44}
45
46impl Default for Cookies {
47    fn default() -> Self {
48        Self {
49            csrf: "".to_string(),
50            session: "".to_string(),
51            site: LeetcodeSite::LeetcodeCom,
52        }
53    }
54}
55
56impl Display for Cookies {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        write!(
59            f,
60            "LEETCODE_SESSION={};csrftoken={};",
61            self.session, self.csrf
62        )
63    }
64}
65
66/// Override cookies from environment variables
67pub const LEETCODE_CSRF_ENV: &str = "LEETCODE_CSRF";
68pub const LEETCODE_SESSION_ENV: &str = "LEETCODE_SESSION";
69pub const LEETCODE_SITE_ENV: &str = "LEETCODE_SITE";
70
71impl Cookies {
72    /// Load cookies from environment variables, overriding any existing values
73    /// if the environment variables are set.
74    pub fn with_env_override(mut self) -> Self {
75        if let Ok(csrf) = std::env::var(LEETCODE_CSRF_ENV) {
76            self.csrf = csrf;
77        }
78        if let Ok(session) = std::env::var(LEETCODE_SESSION_ENV) {
79            self.session = session;
80        }
81        if let Ok(site) = std::env::var(LEETCODE_SITE_ENV) {
82            if let Ok(leetcode_site) = LeetcodeSite::from_str(&site) {
83                self.site = leetcode_site;
84            }
85        }
86        self
87    }
88}