1use std::io::Error;
2use std::path::{Path, PathBuf};
3use std::sync::{LazyLock, OnceLock};
4
5mod config;
6mod db;
7mod ssh;
8
9static WORK_DIR: OnceLock<PathBuf> = OnceLock::new();
10pub(crate) static WORK_DIR_FILE: LazyLock<fn(&str) -> PathBuf> = LazyLock::new(|| {
11 |n| {
12 if cfg!(test) {
13 let work_dir =
14 std::env::current_exe().expect("failed to get current execute directory");
15 work_dir.with_file_name("test.atsh.d")
16 } else {
17 let work_dir = WORK_DIR.get().expect("WORK_DIR not initialized");
18 work_dir.join(n)
19 }
20 }
21});
22
23fn setup_logging(work_dir: &Path) -> Result<(), Error> {
24 use tracing_appender::rolling::{RollingFileAppender, Rotation};
25 use tracing_subscriber::{
26 filter::LevelFilter,
27 fmt::{
28 self,
29 time::{LocalTime, UtcTime},
30 },
31 layer::SubscriberExt,
32 util::SubscriberInitExt,
33 EnvFilter,
34 };
35
36 let log_dir = work_dir.join("logs");
37 if !log_dir.is_dir() {
38 std::fs::create_dir_all(&log_dir)?;
39 }
40
41 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
42 let debug = filter
43 .max_level_hint()
44 .map(|level| level >= LevelFilter::DEBUG)
45 .unwrap_or(false);
46
47 let subscriber = tracing_subscriber::registry()
48 .with(filter)
49 .with(
50 fmt::Layer::new()
51 .with_writer(std::io::stderr)
52 .with_target(debug)
53 .with_line_number(debug)
54 .with_timer(LocalTime::rfc_3339()),
55 )
56 .with(
57 fmt::Layer::new()
58 .json()
59 .with_writer(
60 RollingFileAppender::builder()
61 .rotation(Rotation::DAILY)
62 .filename_suffix("json")
64 .build(log_dir)
65 .expect("Failed to create log file"),
66 )
67 .with_target(debug)
68 .with_line_number(debug)
69 .with_timer(UtcTime::rfc_3339()),
70 );
71
72 subscriber.init();
73 tracing::debug!(
74 "Logging initialized (RUST_LOG={})",
75 std::env::var("RUST_LOG").unwrap_or_default()
76 );
77 Ok(())
78}
79
80pub mod atsh {
81 use crate::{setup_logging, WORK_DIR};
82 use std::io::{Error, ErrorKind};
83 use std::path::Path;
84 use tracing::debug;
85
86 pub use crate::ssh::remote::Remote;
87 use crate::ssh::remote::Remotes;
88
89 type Result<T> = std::result::Result<T, Error>;
90
91 pub fn initialize(work_dir: Option<&Path>) -> Result<()> {
92 let work_dir = work_dir.map(|p| p.to_path_buf()).unwrap_or({
93 let work_dir =
94 std::env::current_exe().expect("failed to get current execute directory");
95 work_dir.with_file_name(".atsh.d")
98 });
99 debug!("work_dir: {}", work_dir.display());
100
101 if !work_dir.exists() {
102 std::fs::create_dir_all(&work_dir)?;
103 }
104 setup_logging(&work_dir)?;
105
106 WORK_DIR
107 .set(work_dir)
108 .expect("WORK_DIR already initialized");
109
110 Ok(())
112 }
113
114 pub fn add(
115 user: &str,
116 password: &str,
117 ip: &str,
118 port: u16,
119 name: &Option<impl AsRef<str>>,
120 note: &Option<impl AsRef<str>>,
121 ) -> Result<usize> {
122 Remotes::add(user, password, ip, port, name, note)
123 }
124
125 pub fn add_remote(remote: &Remote) -> Result<usize> {
126 add(
127 &remote.user,
128 &remote.password,
129 &remote.ip,
130 remote.port,
131 &remote.name,
132 &remote.note,
133 )
134 }
135
136 pub fn remove(index: &Vec<usize>) -> Result<usize> {
137 Remotes::delete(index)
138 }
139
140 pub fn get(index: usize) -> Result<Option<Remote>> {
141 Remotes::get(index)
142 }
143
144 pub fn try_get(index: usize) -> Result<Remote> {
145 Remotes::try_get(index)
146 }
147
148 pub fn get_all() -> Result<Vec<Remote>> {
149 let remotes = Remotes::get_all()?;
150 Ok(remotes.0)
151 }
152
153 pub fn pprint(all: bool) -> Result<()> {
154 let remotes = Remotes::get_all()?;
155 remotes.pprint(all);
156 Ok(())
157 }
158
159 pub fn login(index: usize, auth: bool) -> Result<()> {
169 let remote = Remotes::try_get(index)?;
170 remote.login(auth)
171 }
172
173 pub fn copy(index: usize, path: &str) -> Result<()> {
174 let paths = path.split('=').collect::<Vec<&str>>();
175 if paths.len() != 2 {
176 return Err(Error::new(
177 ErrorKind::InvalidInput,
178 "path format error, like `from=to`",
179 ));
180 }
181 let remote = Remotes::try_get(index)?;
182 if std::path::PathBuf::from(paths[0]).exists() {
183 remote.upload(paths[0], paths[1])
184 } else {
185 remote.download(paths[0], paths[1])
186 }
187 }
188
189 pub fn upload(index: usize, path: &Vec<impl AsRef<str>>) -> Result<()> {
190 if path.len() != 2 {
191 return Err(Error::new(
192 ErrorKind::InvalidInput,
193 "path format error, like `upload -p /local/path /remote/path`",
194 ));
195 }
196 let (local, remote) = (path[0].as_ref(), path[1].as_ref());
197 if !Path::new(local).exists() {
198 return Err(Error::new(ErrorKind::NotFound, "the upload file not found"));
199 }
200
201 Remotes::try_get(index)?.upload(local, remote)
202 }
203
204 pub fn download(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 /remote/path /local/path`",
209 ));
210 }
211 let (remote, local) = (path[0].as_ref(), path[1].as_ref());
212
213 Remotes::try_get(index)?.download(remote, local)
214 }
215}
216
217