atsh_lib/
lib.rs

1use std::env::{current_exe, home_dir};
2use std::io::{Error, ErrorKind};
3use std::path::{Path, PathBuf};
4use std::sync::{LazyLock, OnceLock};
5
6mod config;
7mod db;
8mod ssh;
9
10static WORK_DIR: OnceLock<PathBuf> = OnceLock::new();
11
12fn get_work_dir() -> &'static PathBuf {
13    WORK_DIR.get().expect("WORK_DIR not initialized")
14}
15
16fn set_work_dir(w: Option<impl AsRef<Path>>) -> Result<(), Error> {
17    let work_dir = w
18        // user specified work_dir
19        .map(|d| d.as_ref().to_path_buf())
20        // $HOME/.atsh.d
21        .or_else(|| home_dir().map(|h| h.join(".atsh.d")))
22        // current executable directory
23        .or_else(|| current_exe().ok().map(|e| e.with_file_name(".atsh.d")))
24        // Error
25        .ok_or_else(|| Error::new(ErrorKind::NotFound, "work_dir not found"))?;
26
27    if !work_dir.exists() {
28        std::fs::create_dir_all(&work_dir)?;
29    }
30
31    WORK_DIR
32        .set(work_dir)
33        .expect("WORK_DIR already initialized");
34    Ok(())
35}
36
37pub(crate) static WORK_DIR_FILE: LazyLock<fn(&str) -> PathBuf> = LazyLock::new(|| {
38    |n| {
39        let work_dir = get_work_dir();
40        if cfg!(test) {
41            work_dir.with_file_name("test.atsh.d")
42        } else {
43            work_dir.join(n)
44        }
45    }
46});
47
48fn setup_logging(work_dir: &Path) -> Result<(), Error> {
49    use tracing_appender::rolling::{RollingFileAppender, Rotation};
50    use tracing_subscriber::{
51        filter::LevelFilter,
52        fmt::{
53            self,
54            time::{LocalTime, UtcTime},
55        },
56        layer::SubscriberExt,
57        util::SubscriberInitExt,
58        EnvFilter,
59    };
60
61    let log_dir = work_dir.join("logs");
62    if !log_dir.is_dir() {
63        std::fs::create_dir_all(&log_dir)?;
64    }
65
66    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
67    let debug = filter
68        .max_level_hint()
69        .map(|level| level >= LevelFilter::DEBUG)
70        .unwrap_or(false);
71
72    let subscriber = tracing_subscriber::registry()
73        .with(filter)
74        .with(
75            fmt::Layer::new()
76                .with_writer(std::io::stderr)
77                .with_target(debug)
78                .with_line_number(debug)
79                .with_timer(LocalTime::rfc_3339()),
80        )
81        .with(
82            fmt::Layer::new()
83                .json()
84                .with_writer(
85                    RollingFileAppender::builder()
86                        .rotation(Rotation::DAILY)
87                        // .filename_prefix("atsh.log")
88                        .filename_suffix("json")
89                        .build(log_dir)
90                        .expect("Failed to create log file"),
91                )
92                .with_target(debug)
93                .with_line_number(debug)
94                .with_timer(UtcTime::rfc_3339()),
95        );
96
97    subscriber.init();
98    tracing::debug!(
99        "Logging initialized (RUST_LOG={})",
100        std::env::var("RUST_LOG").unwrap_or_default()
101    );
102    Ok(())
103}
104
105pub mod atsh {
106    use std::io::{Error, ErrorKind};
107    use std::path::Path;
108    use tracing::debug;
109
110    use crate::ssh::remote::Remotes;
111    use crate::{get_work_dir, set_work_dir, setup_logging};
112
113    pub use crate::ssh::{remote::Remote, secure::set_atshkey};
114
115    type Result<T> = std::result::Result<T, Error>;
116
117    pub fn initialize(work_dir: Option<impl AsRef<Path>>) -> Result<()> {
118        set_work_dir(work_dir)?;
119        let w = get_work_dir();
120        setup_logging(w)?;
121        debug!("success initialize at {:?}", w);
122        Ok(())
123    }
124
125    pub fn add(
126        user: &str,
127        password: &str,
128        ip: &str,
129        port: u16,
130        name: &Option<impl AsRef<str>>,
131        note: &Option<impl AsRef<str>>,
132    ) -> Result<usize> {
133        Remotes::add(user, password, ip, port, name, note)
134    }
135
136    pub fn add_remote(remote: &Remote) -> Result<usize> {
137        add(
138            &remote.user,
139            &remote.password,
140            &remote.ip,
141            remote.port,
142            &remote.name,
143            &remote.note,
144        )
145    }
146
147    pub fn remove(index: &Vec<usize>) -> Result<usize> {
148        Remotes::delete(index)
149    }
150
151    pub fn get(index: usize) -> Result<Option<Remote>> {
152        Remotes::get(index)
153    }
154
155    pub fn try_get(index: usize) -> Result<Remote> {
156        Remotes::try_get(index)
157    }
158
159    pub fn get_all() -> Result<Vec<Remote>> {
160        let remotes = Remotes::get_all()?;
161        Ok(remotes.0)
162    }
163
164    pub fn pprint(all: bool) -> Result<()> {
165        let remotes = Remotes::get_all()?;
166        remotes.pprint(all);
167        Ok(())
168    }
169
170    // pub fn list(all: bool) -> Result<()> {
171    //     if all {
172    //         Remotes::list_all()
173    //     } else {
174    //         Remotes::list()
175    //     }
176    // }
177
178    // auth params means try auth against the server
179    pub fn login(index: usize, auth: bool) -> Result<()> {
180        let remote = Remotes::try_get(index)?;
181        remote.login(auth)
182    }
183
184    #[deprecated(
185        since = "0.1.2",
186        note = "This function is not clearly expressed; use `upload/download` instead."
187    )]
188    pub fn copy(index: usize, path: &str) -> Result<()> {
189        let paths = path.split('=').collect::<Vec<&str>>();
190        if paths.len() != 2 {
191            return Err(Error::new(
192                ErrorKind::InvalidInput,
193                "path format error, like `from=to`",
194            ));
195        }
196        let remote = Remotes::try_get(index)?;
197        if std::path::PathBuf::from(paths[0]).exists() {
198            remote.upload(paths[0], paths[1])
199        } else {
200            remote.download(paths[0], paths[1])
201        }
202    }
203
204    pub fn upload(index: usize, path: &Vec<impl AsRef<str>>) -> Result<()> {
205        if path.len() != 2 {
206            return Err(Error::new(
207                ErrorKind::InvalidInput,
208                "path format error, like `upload -p /local/path /remote/path`",
209            ));
210        }
211        let (local, remote) = (path[0].as_ref(), path[1].as_ref());
212        if !Path::new(local).exists() {
213            return Err(Error::new(ErrorKind::NotFound, "the upload file not found"));
214        }
215
216        Remotes::try_get(index)?.upload(local, remote)
217    }
218
219    pub fn download(index: usize, path: &Vec<impl AsRef<str>>) -> Result<()> {
220        if path.len() != 2 {
221            return Err(Error::new(
222                ErrorKind::InvalidInput,
223                "path format error, like `upload -p /remote/path /local/path`",
224            ));
225        }
226        let (remote, local) = (path[0].as_ref(), path[1].as_ref());
227
228        Remotes::try_get(index)?.download(remote, local)
229    }
230}
231
232#[cfg(test)]
233mod tests {
234
235    use super::*;
236
237    #[test]
238    fn test_work_dir_default() {
239        let w = set_work_dir(Option::<&str>::None);
240        assert!(w.is_ok());
241        let w = get_work_dir();
242        println!("default work dir: {:?}", w);
243        if let Some(h) = home_dir() {
244            assert_eq!(w.to_owned(), h.join(".atsh.d"));
245        } else {
246            if let Ok(e) = current_exe() {
247                assert_eq!(w.to_owned(), e.join(".atsh.d"));
248            } else {
249                assert!(false)
250            }
251        }
252    }
253
254    #[test]
255    fn test_work_dir() {
256        let sw = PathBuf::from("test");
257        let w = set_work_dir(Some(&sw));
258        assert!(w.is_ok());
259        let w = get_work_dir();
260        println!("set work dir: {:?}", w);
261        assert_eq!(w.to_owned(), sw);
262    }
263}
atsh_lib/lib.rs

atsh_lib/
lib.rs
