1use std::{collections::HashSet, fmt::Display, fs::File, io::IoSlice, path::PathBuf};
2
3use decrypt_cookies::prelude::*;
4use snafu::ResultExt;
5use strum::IntoEnumIterator;
6use tokio::task;
7
8use crate::{
9 args::{ChromiumName, Value},
10 error::{self, 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 ) -> Result<()>
31 where
32 H: Into<Option<String>>,
33 {
34 let host = host.into();
35 for task in names.map(|name| {
36 let host = host.clone();
37 let output_dir = output_dir.clone();
38 let sep = sep.clone();
39 let values = HashSet::from_iter(Value::iter());
40
41 tokio::task::spawn(async move {
42 Self::write_data(name, None, host, values, output_dir, sep).await
43 })
44 }) {
45 task.await
46 .context(error::TokioTaskSnafu)??;
47 }
48
49 Ok(())
50 }
51
52 pub async fn write_data<D, H, S>(
53 name: ChromiumName,
54 data_dir: D,
55 host: H,
56 values: HashSet<Value>,
57 mut output_dir: PathBuf,
58 sep: S,
59 ) -> Result<()>
60 where
61 D: Into<Option<PathBuf>>,
62 H: Into<Option<String>>,
63 S: Display + Send + Clone + 'static,
64 {
65 let data_dir = data_dir.into();
66 let host: Option<String> = host.into();
67
68 macro_rules! chromiums {
69 ($($browser:ident,) *) => {
70 match name {
71 $(
72 ChromiumName::$browser => {
73 let chromium = if let Some(dir) = data_dir {
74 ChromiumBuilder::<$browser>::with_user_data_dir(dir)
75 }
76 else {
77 ChromiumBuilder::new()
78 }
79 .build()
80 .await
81 .context(error::ChromiumBuilderSnafu)?;
82
83 let cookies = if values.contains(&Value::Cookie) {
84 let host = host.clone();
85 let chromium = chromium.clone();
86 let task = task::spawn(async move {
87 let cookies = if let Some(host) = host {
88 chromium
89 .cookies_by_host(host)
90 .await
91 }
92 else {
93 chromium.cookies_all().await
94 }
95 .context(error::ChromiumSnafu)?;
96 Ok::<_, error::Error>(cookies)
97 });
98 Some(task)
99 }
100 else {
101 None
102 };
103
104 let logins = if values.contains(&Value::Login) {
105 let host = host.clone();
106 let task = task::spawn(async move {
107 let logins = if let Some(host) = host {
108 chromium.logins_by_host(host).await
109 }
110 else {
111 chromium.all_logins().await
112 }
113 .context(error::ChromiumSnafu)?;
114 Ok::<_, error::Error>(logins)
115 });
116 Some(task)
117 }
118 else {
119 None
120 };
121 (cookies, logins, $browser::NAME)
122 },
123 )*
124 }
125 };
126 }
127
128 #[cfg(target_os = "linux")]
129 let (cookies, logins, name) =
130 chromiums![Chrome, Edge, Chromium, Brave, Vivaldi, Yandex, Opera,];
131 #[cfg(not(target_os = "linux"))]
132 let (cookies, logins, name) = chromiums![
133 Chrome, Edge, Chromium, Brave, Vivaldi, Yandex, Opera, Arc, OperaGX, CocCoc,
134 ];
135 let (cookies, logins, cap) = match (cookies, logins) {
136 (None, None) => (None, None, 0),
137 (None, Some(logins)) => {
138 let l = logins
139 .await
140 .context(error::TokioTaskSnafu)??;
141 (None, Some(l), 1)
142 },
143 (Some(cookies), None) => {
144 let c = cookies
145 .await
146 .context(error::TokioTaskSnafu)??;
147 (Some(c), None, 1)
148 },
149 (Some(cookies), Some(logins)) => {
150 let (c, l) = tokio::join!(cookies, logins);
151 let c = c.context(error::TokioTaskSnafu)??;
152 let l = l.context(error::TokioTaskSnafu)??;
153 (Some(c), Some(l), 2)
154 },
155 };
156
157 output_dir.push(name);
158 tokio::fs::create_dir_all(&output_dir)
159 .await
160 .context(error::IoSnafu { path: output_dir.clone() })?;
161
162 let mut tasks = Vec::with_capacity(cap);
163
164 if let Some(cookies) = cookies {
165 let out_file = output_dir.join(crate::COOKIES_FILE);
166 let sep = sep.clone();
167
168 let handle = utils::write_cookies(out_file, cookies, sep);
169 tasks.push(handle);
170 }
171
172 if let Some(logins) = logins {
173 let out_file = output_dir.join(crate::LOGINS_FILE);
174 let sep = sep.clone();
175
176 let handle = task::spawn_blocking(move || {
177 let mut file = File::options()
178 .write(true)
179 .create(true)
180 .truncate(true)
181 .open(&out_file)
182 .context(error::IoSnafu { path: out_file.clone() })?;
183
184 let header = login_csv_header(sep.clone());
185
186 let mut slices = Vec::with_capacity(2 + logins.len() * 2);
187 slices.push(IoSlice::new(header.as_bytes()));
188 slices.push(IoSlice::new(b"\n"));
189
190 let csvs: Vec<_> = logins
191 .into_iter()
192 .map(|v| v.to_csv(sep.clone()))
193 .collect();
194
195 for csv in &csvs {
196 slices.push(IoSlice::new(csv.as_bytes()));
197 slices.push(IoSlice::new(b"\n"));
198 }
199 write_all_vectored(&mut file, &mut slices)
200 .context(error::IoSnafu { path: out_file.clone() })?;
201 Ok::<(), error::Error>(())
202 });
203 tasks.push(handle);
204 }
205
206 for ele in tasks {
207 ele.await
208 .context(error::TokioTaskSnafu)??;
209 }
210
211 Ok(())
212 }
213}