Skip to main content

tradingview/
user.rs

1pub use crate::models::UserCookies;
2use crate::{
3    Result,
4    error::{Error, LoginError},
5    utils::http_client,
6};
7use google_authenticator::{GA_AUTH, get_code};
8use reqwest::{
9    Response,
10    header::{CONTENT_TYPE, COOKIE},
11};
12use serde::Deserialize;
13use serde_json::Value;
14use tracing::{debug, error, info, warn};
15
16impl UserCookies {
17    pub fn new() -> Self {
18        Default::default()
19    }
20
21    pub async fn login(
22        &mut self,
23        username: &str,
24        password: &str,
25        totp_secret: Option<&str>,
26    ) -> Result<Self> {
27        let client = http_client();
28        let response = client
29            .post("https://www.tradingview.com/accounts/signin/")
30            .header(CONTENT_TYPE, "application/x-www-form-urlencoded")
31            .body(format!(
32                "username={username}&password={password}&remember=true"
33            ))
34            .send()
35            .await?;
36
37        let (session, signature, device_token) =
38            response
39                .cookies()
40                .fold((None, None, None), |session_cookies, cookie| {
41                    match cookie.name() {
42                        "sessionid" => (
43                            Some(cookie.value().to_string()),
44                            session_cookies.1,
45                            session_cookies.2,
46                        ),
47                        "sessionid_sign" => (
48                            session_cookies.0,
49                            Some(cookie.value().to_string()),
50                            session_cookies.2,
51                        ),
52                        "device_t" => (
53                            session_cookies.0,
54                            session_cookies.1,
55                            Some(cookie.value().to_string()),
56                        ),
57                        _ => session_cookies,
58                    }
59                });
60        if session.is_none() || signature.is_none() {
61            error!("unable to login, username or password is invalid");
62            return Err(Error::Login {
63                source: LoginError::InvalidCredentials,
64            });
65        }
66
67        #[derive(Debug, Deserialize)]
68        struct LoginUserResponse {
69            user: UserCookies,
70        }
71
72        let response: Value = response.json().await?;
73
74        if response["error"] == *"" {
75            debug!("User data: {:#?}", response);
76            warn!("2FA is not enabled for this account");
77            info!("User is logged in successfully");
78            let login_resp: LoginUserResponse = serde_json::from_value(response)?;
79
80            Ok(UserCookies {
81                session: session.unwrap_or_default(),
82                session_signature: signature.unwrap_or_default(),
83                device_token: device_token.unwrap_or_default(),
84                ..login_resp.user
85            })
86        } else if response["error"] == *"2FA_required" {
87            if totp_secret.is_none() {
88                error!("2FA is enabled for this account, but no TOTP secret was provided");
89                return Err(Error::Login {
90                    source: LoginError::OTPSecretNotFound,
91                });
92            }
93
94            let response = Self::handle_mfa(
95                totp_secret.unwrap(),
96                session.clone().unwrap_or_default().as_str(),
97                signature.clone().unwrap_or_default().as_str(),
98            )
99            .await?;
100
101            let (session, signature, device_token) =
102                response
103                    .cookies()
104                    .fold((None, None, None), |session_cookies, cookie| {
105                        match cookie.name() {
106                            "sessionid" => (
107                                Some(cookie.value().to_string()),
108                                session_cookies.1,
109                                session_cookies.2,
110                            ),
111                            "sessionid_sign" => (
112                                session_cookies.0,
113                                Some(cookie.value().to_string()),
114                                session_cookies.2,
115                            ),
116                            "device_t" => (
117                                session_cookies.0,
118                                session_cookies.1,
119                                Some(cookie.value().to_string()),
120                            ),
121                            _ => session_cookies,
122                        }
123                    });
124
125            let body = response.json().await?;
126            debug!("2FA login response: {:#?}", body);
127            info!("User is logged in successfully");
128            let login_resp: LoginUserResponse = serde_json::from_value(body)?;
129
130            Ok(UserCookies {
131                session: session.unwrap_or_default(),
132                session_signature: signature.unwrap_or_default(),
133                device_token: device_token.unwrap_or_default(),
134                ..login_resp.user
135            })
136        } else {
137            error!("unable to login, username or password is invalid");
138            Err(Error::Login {
139                source: LoginError::InvalidCredentials,
140            })
141        }
142    }
143
144    async fn handle_mfa(totp_secret: &str, session: &str, signature: &str) -> Result<Response> {
145        if totp_secret.is_empty() {
146            return Err(Error::Login {
147                source: LoginError::OTPSecretNotFound,
148            });
149        }
150
151        let cookie = format!("sessionid={session}; sessionid_sign={signature};");
152        let response = http_client()
153            .post("https://www.tradingview.com/accounts/two-factor/signin/totp/")
154            .header(COOKIE, &cookie)
155            .header(CONTENT_TYPE, "application/x-www-form-urlencoded")
156            .body(format!(
157                "code={}",
158                match get_code!(totp_secret) {
159                    Ok(code) => code,
160                    Err(e) => {
161                        error!("Error generating TOTP code: {}", e);
162                        return Err(Error::Login {
163                            source: LoginError::InvalidOTPSecret,
164                        });
165                    }
166                }
167            ))
168            .send()
169            .await?;
170
171        if response.status().is_success() {
172            Ok(response)
173        } else {
174            Err(Error::Login {
175                source: LoginError::InvalidOTPSecret,
176            })
177        }
178    }
179}