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