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}