1use std::{collections::HashSet, fmt::Display, fs::File, io::IoSlice, path::PathBuf};
2
3use decrypt_cookies::{
4 chromium::{builder::ChromiumBuilderError, GetCookies, GetLogins},
5 prelude::*,
6};
7use snafu::ResultExt;
8use strum::IntoEnumIterator;
9use tokio::task;
10
11use crate::{
12 args::{ChromiumName, Format, Value},
13 error::{self, IoSnafu, JsonSnafu, Result},
14 utils::{self, write_all_vectored},
15};
16
17#[derive(Clone, Copy)]
18#[derive(Debug)]
19#[derive(Default)]
20#[derive(PartialEq, Eq, PartialOrd, Ord)]
21pub struct ChromiumBased;
22
23fn login_csv_header<D: Display>(sep: D) -> String {
24 format!("url{sep}username{sep}display_name{sep}password{sep}date_created{sep}date_last_used{sep}modified")
25}
26
27impl ChromiumBased {
28 pub(crate) async fn multi_data<H>(
29 names: impl Iterator<Item = ChromiumName>,
30 output_dir: PathBuf,
31 sep: String,
32 host: H,
33 format: Format,
34 ) -> Result<()>
35 where
36 H: Into<Option<String>>,
37 {
38 let host = host.into();
39 for task in names.map(|name| {
40 let host = host.clone();
41 let output_dir = output_dir.clone();
42 let sep = sep.clone();
43 let values = HashSet::from_iter(Value::iter());
44
45 tokio::task::spawn(async move {
46 Self::write_data(name, None, host, values, output_dir, sep, format).await
47 })
48 }) {
49 if let Err(e) = task
50 .await
51 .context(error::TokioTaskSnafu)?
52 {
53 match e {
54 error::Error::ChromiumBuilder {
55 source: source @ ChromiumBuilderError::NotFoundBase { .. },
56 ..
57 } => {
58 #[cfg(not(target_os = "windows"))]
59 tracing::info!(r#"{source}"#,);
60 #[cfg(target_os = "windows")]
61 tracing::info!(
62 r#"{source}
63When you use scoop on Windows, the data path is located at `~\scoop\persisst\<name>\<xxx>`"#,
64 );
65 },
66 e => tracing::error!("{e}"),
67 }
68 }
69 }
70
71 Ok(())
72 }
73
74 pub async fn write_data<D, H, S>(
75 name: ChromiumName,
76 data_dir: D,
77 host: H,
78 values: HashSet<Value>,
79 mut output_dir: PathBuf,
80 sep: S,
81 format: Format,
82 ) -> Result<()>
83 where
84 D: Into<Option<PathBuf>>,
85 H: Into<Option<String>>,
86 S: Display + Send + Clone + 'static,
87 {
88 let data_dir = data_dir.into();
89 let host: Option<String> = host.into();
90
91 macro_rules! chromiums {
92 ($($browser:ident,) *) => {
93 match name {
94 $(
95 ChromiumName::$browser => {
96 let chromium = if let Some(dir) = data_dir {
97 ChromiumBuilder::<$browser>::with_user_data_dir(dir)
98 }
99 else {
100 ChromiumBuilder::new()
101 }
102 .build()
103 .await
104 .context(error::ChromiumBuilderSnafu)?;
105
106 let cookies = if values.contains(&Value::Cookie) {
107 let host = host.clone();
108 let chromium = chromium.clone();
109 let task = task::spawn(async move {
110 let cookies = if let Some(host) = host {
111 chromium
112 .cookies_by_host(host)
113 .await
114 }
115 else {
116 chromium.cookies_all().await
117 }
118 .context(error::ChromiumSnafu)?;
119 Ok::<_, error::Error>(cookies)
120 });
121 Some(task)
122 }
123 else {
124 None
125 };
126
127 let logins = if values.contains(&Value::Login) {
128 let host = host.clone();
129 let task = task::spawn(async move {
130 let logins = if let Some(host) = host {
131 chromium.logins_by_host(host).await
132 }
133 else {
134 chromium.logins_all().await
135 }
136 .context(error::ChromiumSnafu)?;
137 Ok::<_, error::Error>(logins)
138 });
139 Some(task)
140 }
141 else {
142 None
143 };
144 (cookies, logins, $browser::NAME)
145 },
146 )*
147 }
148 };
149 }
150
151 #[cfg(target_os = "linux")]
152 let (cookies, logins, name) =
153 chromiums![Chrome, Edge, Chromium, Brave, Vivaldi, Yandex, Opera,];
154 #[cfg(not(target_os = "linux"))]
155 let (cookies, logins, name) = chromiums![
156 Chrome, Edge, Chromium, Brave, Vivaldi, Yandex, Opera, Arc, OperaGX, CocCoc,
157 ];
158 let (cookies, logins, cap) = match (cookies, logins) {
159 (None, None) => (None, None, 0),
160 (None, Some(logins)) => {
161 let l = logins
162 .await
163 .context(error::TokioTaskSnafu)??;
164 (None, Some(l), 1)
165 },
166 (Some(cookies), None) => {
167 let c = cookies
168 .await
169 .context(error::TokioTaskSnafu)??;
170 (Some(c), None, 1)
171 },
172 (Some(cookies), Some(logins)) => {
173 let (c, l) = tokio::join!(cookies, logins);
174 let c = c.context(error::TokioTaskSnafu)??;
175 let l = l.context(error::TokioTaskSnafu)??;
176 (Some(c), Some(l), 2)
177 },
178 };
179
180 output_dir.push(name);
181 tokio::fs::create_dir_all(&output_dir)
182 .await
183 .with_context(|_| error::IoSnafu { path: output_dir.clone() })?;
184
185 let mut tasks = Vec::with_capacity(cap);
186
187 if let Some(cookies) = cookies {
188 let out_file = output_dir.join({
189 match format {
190 Format::Csv => crate::COOKIES_FILE_CSV,
191 Format::Json => crate::COOKIES_FILE_JSON,
192 Format::JsonLines => crate::COOKIES_FILE_JSONL,
193 }
194 });
195 let sep = sep.clone();
196
197 let handle = utils::write_cookies(out_file, cookies, sep, format);
198 tasks.push(handle);
199 }
200
201 if let Some(logins) = logins {
202 let sep = sep.clone();
203
204 let handle = task::spawn_blocking(move || {
205 let out_file = output_dir.join(match format {
206 Format::Csv => crate::LOGINS_FILE_CSV,
207 Format::Json => crate::LOGINS_FILE_JSON,
208 Format::JsonLines => crate::LOGINS_FILE_JSONL,
209 });
210
211 let mut file = File::options()
212 .write(true)
213 .create(true)
214 .truncate(true)
215 .open(&out_file)
216 .with_context(|_| error::IoSnafu { path: out_file.clone() })?;
217
218 match format {
219 Format::Csv => {
220 let mut slices = Vec::with_capacity(2 + logins.len() * 2);
221
222 let header = login_csv_header(sep.clone());
223 slices.push(IoSlice::new(header.as_bytes()));
224 slices.push(IoSlice::new(b"\n"));
225
226 let csvs: Vec<_> = logins
227 .into_iter()
228 .map(|v| v.to_csv(sep.clone()))
229 .collect();
230
231 for csv in &csvs {
232 slices.push(IoSlice::new(csv.as_bytes()));
233 slices.push(IoSlice::new(b"\n"));
234 }
235 write_all_vectored(&mut file, &mut slices)
236 .with_context(|_| error::IoSnafu { path: out_file })
237 },
238 Format::Json => serde_json::to_writer(file, &logins).context(JsonSnafu),
239 Format::JsonLines => serde_jsonlines::write_json_lines(&out_file, &logins)
240 .context(IoSnafu { path: out_file }),
241 }
242 });
243 tasks.push(handle);
244 }
245
246 for ele in tasks {
247 ele.await
248 .context(error::TokioTaskSnafu)??;
249 }
250
251 Ok(())
252 }
253}