ibdl_core/cli/
mod.rs

1// 20002709
2use ibdl_common::post::{extension::Extension, NameType};
3use ibdl_extractors::extractor_config::ServerConfig;
4use once_cell::sync::OnceCell;
5use std::{collections::HashMap, path::PathBuf};
6
7use clap::{Parser, Subcommand};
8
9use crate::generate_output_path_precise;
10
11use self::{
12    commands::{pool::Pool, post::Post, search::TagSearch},
13    extra::validate_imageboard,
14};
15
16pub mod commands;
17pub(crate) mod extra;
18
19pub static AVAILABLE_SERVERS: OnceCell<HashMap<String, ServerConfig>> = OnceCell::new();
20
21#[derive(Debug, Subcommand)]
22pub enum Commands {
23    /// Search and download posts with tags
24    Search(TagSearch),
25    /// Download entire pools of posts
26    Pool(Pool),
27    /// Download a single or multiple specific posts
28    Post(Post),
29}
30
31#[derive(Parser, Debug)]
32#[clap(name = "Imageboard Downloader", author, version, about, long_about = None)]
33pub struct Cli {
34    #[clap(subcommand)]
35    pub mode: Commands,
36
37    /// Specify which website to download from
38    ///
39    /// Default websites include: ["danbooru", "e621", "gelbooru", "rule34", "realbooru", "konachan"]
40    #[clap(short, long, ignore_case = true, default_value_t = ServerConfig::default(), global = true, value_parser = validate_imageboard)]
41    pub imageboard: ServerConfig,
42
43    /// Print all available servers and exit
44    #[clap(long, global = true)]
45    pub servers: bool,
46
47    /// Where to save files (If the path doesn't exist, it will be created.)
48    #[clap(short = 'o', value_name = "PATH", help_heading = "SAVE", global = true)]
49    pub output: Option<PathBuf>,
50
51    /// Number of simultaneous downloads
52    ///
53    /// [max: 20]
54    #[clap(
55        short = 'd',
56        value_name = "NUMBER",
57        global = true,
58        value_parser(clap::value_parser!(u8).range(1..=20)),
59        default_value_t = 5,
60        help_heading = "DOWNLOAD",
61        global = true,
62    )]
63    pub simultaneous_downloads: u8,
64
65    /// Authenticate to the imageboard website.
66    ///
67    /// This flag only needs to be set a single time.
68    ///
69    /// Once authenticated, it's possible to use your blacklist to exclude posts with unwanted tags
70    #[clap(short, long, action, help_heading = "GENERAL", global = true)]
71    pub auth: bool,
72
73    /// Save files with their ID as filename instead of it's MD5
74    ///
75    /// If the output dir has the same file downloaded with the MD5 name, it will be renamed to the post's ID
76    #[clap(
77        long = "id",
78        value_parser,
79        global = true,
80        default_value_t = false,
81        help_heading = "SAVE"
82    )]
83    pub save_file_as_id: bool,
84
85    /// Save posts inside a cbz file.
86    ///
87    /// Will ask to overwrite the destination file.
88    #[clap(
89        long,
90        value_parser,
91        default_value_t = false,
92        help_heading = "SAVE",
93        global = true
94    )]
95    pub cbz: bool,
96
97    /// Write tags in a txt file next to the downloaded image (for Stable Diffusion training)
98    #[clap(
99        long,
100        value_parser,
101        default_value_t = false,
102        help_heading = "SAVE",
103        global = true
104    )]
105    pub annotate: bool,
106
107    /// Always overwrite output
108    #[clap(
109        short = 'y',
110        value_parser,
111        default_value_t = false,
112        help_heading = "SAVE",
113        global = true
114    )]
115    pub overwrite: bool,
116}
117
118impl Cli {
119    pub const fn name_type(&self) -> NameType {
120        if self.save_file_as_id {
121            NameType::ID
122        } else {
123            NameType::MD5
124        }
125    }
126
127    pub fn get_extension(&self) -> Option<Extension> {
128        match &self.mode {
129            Commands::Search(args) => {
130                if let Some(ext) = &args.force_extension {
131                    return Some(Extension::guess_format(ext));
132                }
133            }
134            Commands::Pool(args) => {
135                if let Some(ext) = &args.force_extension {
136                    return Some(Extension::guess_format(ext));
137                }
138            }
139            Commands::Post(_) => {}
140        }
141        None
142    }
143
144    pub fn generate_save_path(&self) -> Result<PathBuf, std::io::Error> {
145        let raw_save_path = if let Some(precise_path) = &self.output {
146            precise_path.to_owned()
147        } else {
148            std::env::current_dir()?
149        };
150
151        let dirname = if self.output.is_some() {
152            generate_output_path_precise(&raw_save_path, self.cbz)
153        } else {
154            raw_save_path
155        };
156
157        Ok(dirname)
158    }
159}
ibdl_core/cli/mod.rs

ibdl_core/cli/
mod.rs
