tidy_browser/
utils.rs

1use std::{
2    fmt::Display,
3    fs::File,
4    io::{ErrorKind, IoSlice, Write},
5    path::PathBuf,
6};
7
8use decrypt_cookies::{chromium::ChromiumCookie, prelude::cookies::CookiesInfo};
9use serde::Serialize;
10use snafu::ResultExt;
11use tokio::task::{self, JoinHandle};
12
13use crate::{
14    args::Format,
15    error::{self, IoSnafu, JsonSnafu},
16};
17
18/// Copy from nightly [`Write::write_all_vectored`]
19/// TODO: Use std method
20pub(crate) fn write_all_vectored(
21    file: &mut File,
22    mut bufs: &mut [IoSlice<'_>],
23) -> std::io::Result<()> {
24    IoSlice::advance_slices(&mut bufs, 0);
25    while !bufs.is_empty() {
26        match file.write_vectored(bufs) {
27            Ok(0) => {
28                return Err(std::io::Error::new(
29                    std::io::ErrorKind::WriteZero,
30                    "failed to write whole buffer",
31                ));
32            },
33            Ok(n) => IoSlice::advance_slices(&mut bufs, n),
34            Err(ref e) if matches!(e.kind(), ErrorKind::Interrupted) => {},
35            Err(e) => return Err(e),
36        }
37    }
38    Ok(())
39}
40
41pub(crate) fn write_cookies<S>(
42    out_file: PathBuf,
43    cookies: Vec<impl CookiesInfo + Send + Serialize + 'static>,
44    sep: S,
45    format: Format,
46) -> JoinHandle<Result<(), error::Error>>
47where
48    S: Display + Send + Clone + 'static,
49{
50    task::spawn_blocking(move || {
51        let mut file = File::options()
52            .write(true)
53            .create(true)
54            .truncate(true)
55            .open(&out_file)
56            .with_context(|_| error::IoSnafu { path: out_file.clone() })?;
57
58        match format {
59            Format::Csv => {
60                let mut slices = Vec::with_capacity(2 + cookies.len() * 2);
61
62                let header = <ChromiumCookie as CookiesInfo>::csv_header(sep.clone());
63                slices.push(IoSlice::new(header.as_bytes()));
64                slices.push(IoSlice::new(b"\n"));
65
66                let csvs: Vec<_> = cookies
67                    .into_iter()
68                    .map(|v| v.to_csv(sep.clone()))
69                    .collect();
70
71                for csv in &csvs {
72                    slices.push(IoSlice::new(csv.as_bytes()));
73                    slices.push(IoSlice::new(b"\n"));
74                }
75
76                write_all_vectored(&mut file, &mut slices)
77                    .with_context(|_| error::IoSnafu { path: out_file })
78            },
79            Format::Json => serde_json::to_writer(file, &cookies).context(JsonSnafu),
80            Format::JsonLines => serde_jsonlines::write_json_lines(&out_file, &cookies)
81                .context(IoSnafu { path: out_file }),
82        }
83    })
84}
tidy_browser/utils.rs

tidy_browser/
utils.rs
