atsh_lib/connection/
remote.rs

1use prettytable::{Cell, Row, Table};
2use serde::{Deserialize, Deserializer, Serialize, Serializer};
3// use std::io::{BufRead, BufReader};
4use std::io::Error;
5use tokio::io::{AsyncBufReadExt, BufReader};
6use tokio::net::TcpStream;
7use tokio::process::Command;
8use tokio::time::Duration;
9use tracing::{debug, info, warn};
10
11use super::rssh::SSHSession;
12use crate::config::CONFIG;
13use crate::storage::db::{
14    delete_index, get_connection, insert, query_all, query_index, update_authorized,
15};
16use crate::storage::secure::{decrypt, encrypt};
17use crate::Result;
18
19#[derive(Debug, Default, Serialize, Deserialize)]
20pub struct Remote {
21    /// the index of the remote server.
22    // #[serde(rename = "idx")]
23    pub index: usize,
24    /// the login user.
25    pub user: String,
26    /// the login password.
27    #[serde(deserialize_with = "depass", serialize_with = "enpass")]
28    pub password: String,
29    /// the login id address.
30    pub ip: String,
31    /// the login port.
32    pub port: u16,
33    /// have the authorized to login.
34    pub authorized: bool,
35    /// the alias name for the login.
36    pub name: Option<String>,
37    /// the note for the server.
38    pub note: Option<String>,
39}
40
41fn depass<'de, D>(deserializer: D) -> std::result::Result<String, D::Error>
42where
43    D: Deserializer<'de>,
44{
45    let password = String::deserialize(deserializer)?;
46    Ok(decrypt(&password))
47}
48
49fn enpass<S>(password: &String, serializer: S) -> std::result::Result<S::Ok, S::Error>
50where
51    S: Serializer,
52{
53    Serialize::serialize(&encrypt(password), serializer)
54}
55
56// impl display for Remote
57impl std::fmt::Display for Remote {
58    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
59        write!(f, "{}@{}:{}", self.user, self.ip, self.port,)
60    }
61}
62
63impl Remote {
64    pub fn add_record(&self) -> Result<usize> {
65        // Force check the ATSH_KEY exist or not
66        CONFIG.get_enc_key()?;
67        let n = {
68            let conn = get_connection().lock();
69            insert(&conn, &self)
70        }
71        .map_err(|e| Error::new(std::io::ErrorKind::Other, e))?;
72        info!(remote = self.to_string(), "success add record");
73        Ok(n)
74    }
75
76    pub fn delete_record(&self) -> Result<()> {
77        // 删除数据库
78        let conn = get_connection().lock();
79        delete_index(&conn, self.index).map_err(|e| Error::new(std::io::ErrorKind::Other, e))?;
80        info!(remote = self.to_string(), "success delete record");
81        Ok(())
82    }
83
84    pub async fn add_auth(&self) -> Result<()> {
85        // check the `ATSH_KEY` exist or not
86        // if not exist, { kind: Other, error: "Authentication failed (username/password)" }
87        let _ = CONFIG.get_enc_key()?;
88        let mut session = SSHSession::new(&self.user, &self.password, &self.ip, self.port).await?;
89        session.authenticate(CONFIG.read_public()?.as_str()).await?;
90        // 更新数据库
91        if !self.authorized {
92            // update authorized to database
93            // self.authorized = true;
94            let conn = get_connection().lock();
95            update_authorized(&conn, self.index, true)?;
96        }
97        info!(remote = self.to_string(), "success add authenticate");
98        Ok(())
99    }
100
101    pub async fn remove_auth(&self) -> Result<()> {
102        let mut session = SSHSession::new(&self.user, &self.password, &self.ip, self.port).await?;
103        session.revoke(CONFIG.read_public()?.as_str()).await?;
104        info!(remote = self.to_string(), "success remove authenticate");
105        Ok(())
106    }
107
108    pub async fn login(&self, reauth: bool) -> Result<()> {
109        // 如果没有认证，或者通过 `--auth` 参数重新认证
110        if !self.authorized || reauth {
111            self.add_auth().await?;
112        }
113        let sshkey = CONFIG.get_private();
114        let _ = Command::new("ssh")
115            .arg(format!("{}@{}", self.user, self.ip))
116            .arg("-p")
117            .arg(self.port.to_string())
118            .arg("-i")
119            .arg(sshkey.to_str().unwrap())
120            .status()
121            .await?;
122        info!(remote = self.to_string(), "success login");
123        Ok(())
124    }
125
126    async fn scp(&self, args: &Vec<&str>) -> Result<()> {
127        // 如果没有认证，则先认证
128        if !self.authorized {
129            debug!(remote = self.to_string(), "no authorized, try authenticate");
130            self.add_auth().await?;
131        }
132        // info!("\n🚨 scp {}\n🚨 input `y` to run and other to cancel.", cmd);
133        // let mut read = String::new();
134        // std::io::stdin().read_line(&mut read)?;
135        // let read = read.trim();
136        // if read == "y" {}
137        debug!(args=?args, "scp");
138
139        let mut child = Command::new("scp")
140            .args(args)
141            .stderr(std::process::Stdio::piped())
142            .spawn()?;
143
144        let stderr = child
145            .stderr
146            .take()
147            .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stderr is none"))?;
148
149        let reader = BufReader::new(stderr);
150        let mut lines = reader.lines();
151
152        while let Some(line) = lines.next_line().await? {
153            println!("{}", line);
154        }
155
156        let _ = child.wait().await?;
157
158        Ok(())
159    }
160
161    pub async fn upload(&self, from: &str, to: &str) -> Result<()> {
162        // debug!(from = ?from, to=?to, "upload");
163        // scp -r -P 22 -i /home/idhyt/.ssh/id_rsa ./test.txt idhyt@1.2.3.4:/tmp
164        let port = self.port.to_string();
165        let remote = format!("{}@{}:{}", self.user, self.ip, to);
166        let cmd = vec![
167            "-r",
168            "-P",
169            &port,
170            "-i",
171            CONFIG.get_private().to_str().unwrap(),
172            from,
173            &remote,
174        ];
175        self.scp(&cmd).await?;
176        info!(from=?from, to=?to, "susccess upload");
177        Ok(())
178    }
179
180    pub async fn download(&self, from: &str, to: &str) -> Result<()> {
181        // debug!(from = ?from, to=?to, "download");
182        // scp -r -P 22 -i /home/idhyt/.ssh/id_rsa idhyt@1.2.3.4:/tmp/test.txt ./
183        let port = self.port.to_string();
184        let remote = format!("{}@{}:{}", self.user, self.ip, from);
185        let cmd = vec![
186            "-r",
187            "-P",
188            &port,
189            "-i",
190            CONFIG.get_private().to_str().unwrap(),
191            &remote,
192            to,
193        ];
194        self.scp(&cmd).await?;
195        info!(from=?from, to=?to, "susccess download");
196        Ok(())
197    }
198
199    async fn is_alive(&self) -> bool {
200        check_port(&self.ip, self.port, None).await
201    }
202}
203
204async fn check_port(host: &str, port: u16, timeout_duration: Option<Duration>) -> bool {
205    let addr = format!("{}:{}", host, port);
206    match tokio::time::timeout(
207        timeout_duration.unwrap_or(Duration::from_secs(3)),
208        TcpStream::connect(&addr),
209    )
210    .await
211    {
212        Ok(Ok(_stream)) => {
213            debug!("✅ {addr} is OPEN");
214            true
215        }
216        Ok(Err(e)) => {
217            warn!("❌ {addr} is CLOSED: {e}");
218            false
219        }
220        Err(e) => {
221            warn!("⏱️ {addr} TIMEOUT {e}");
222            false
223        }
224    }
225}
226
227// #[derive(Serialize, Deserialize)]
228pub struct Remotes(pub Vec<Remote>);
229
230impl Remotes {
231    fn load() -> Result<Remotes> {
232        let conn = get_connection().lock();
233        let remotes = query_all(&conn).map_err(|e| Error::new(std::io::ErrorKind::Other, e))?;
234        Ok(Remotes(remotes))
235    }
236    pub fn get(idx: usize) -> Result<Option<Remote>> {
237        let remote = {
238            let conn = get_connection().lock();
239            query_index(&conn, idx)
240        }
241        .map_err(|e| Error::new(std::io::ErrorKind::Other, e))?;
242
243        if remote.is_some() {
244            info!(index = idx, "susccess get remote");
245        } else {
246            warn!(index = idx, "remote not found");
247        }
248        Ok(remote)
249    }
250
251    pub fn try_get(idx: usize) -> Result<Remote> {
252        if let Some(remote) = Remotes::get(idx)? {
253            Ok(remote)
254        } else {
255            Err(Box::new(Error::new(
256                std::io::ErrorKind::NotFound,
257                format!("index {} remote not found", idx),
258            )))
259        }
260    }
261
262    pub fn get_all() -> Result<Remotes> {
263        let remotes = Remotes::load()?;
264        info!("susccess get all remotes {}", remotes.0.len());
265        Ok(remotes)
266    }
267
268    pub fn add(
269        user: &str,
270        password: &str,
271        ip: &str,
272        port: u16,
273        name: &Option<impl AsRef<str>>,
274        note: &Option<impl AsRef<str>>,
275    ) -> Result<usize> {
276        let remote = Remote {
277            index: 0, // not used
278            user: user.to_string(),
279            password: password.to_string(),
280            ip: ip.to_string(),
281            port,
282            authorized: false,
283            name: name.as_ref().map(|n| n.as_ref().to_string()),
284            note: note.as_ref().map(|n| n.as_ref().to_string()),
285        };
286        // we not authorized the remote server until the first login
287        // remote.authorized();
288        // debug!(remote = remote.to_string(), "add");
289        remote.add_record()
290    }
291
292    pub async fn delete(indexs: &Vec<usize>) -> Result<usize> {
293        let remotes: Vec<Remote> = indexs
294            .iter()
295            .collect::<std::collections::HashSet<_>>()
296            .iter()
297            .filter_map(|&idx| Remotes::get(*idx).transpose())
298            .collect::<std::result::Result<Vec<_>, _>>()?;
299        debug!(total = remotes.len(), "delete");
300        for remote in remotes.iter() {
301            debug!(index = remote.index, "delete");
302            // remove auth
303            if remote.authorized && remote.is_alive().await {
304                remote.remove_auth().await?;
305            }
306            // delete database
307            remote.delete_record()?;
308        }
309        Ok(remotes.len())
310    }
311
312    pub fn pprint(&self, all: bool) {
313        let mut table = Table::new();
314        let mut titles = vec!["index", "name", "user", "ip", "port"];
315        if all {
316            titles.push("password");
317            titles.push("authorized");
318            titles.push("note");
319        }
320        table.set_titles(Row::new(
321            titles
322                .iter()
323                .map(|v| Cell::new(v).style_spec("bcFg"))
324                .collect::<Vec<Cell>>(),
325        ));
326
327        for remote in self.0.iter() {
328            let mut row = vec![
329                remote.index.to_string(),
330                remote.name.clone().unwrap_or_else(|| "".to_string()),
331                remote.user.clone(),
332                remote.ip.clone(),
333                remote.port.to_string(),
334            ];
335            if all {
336                if let Ok(_) = CONFIG.get_enc_key() {
337                    row.push(remote.password.clone());
338                } else {
339                    row.push(format!(
340                        "{}..{}",
341                        &remote.password[..3],
342                        &remote.password[remote.password.len() - 5..]
343                    ));
344                }
345                // row.push(remote.password.clone());
346                row.push(remote.authorized.to_string());
347                row.push(remote.note.clone().unwrap_or_else(|| "".to_string()));
348            }
349            table.add_row(Row::new(
350                row.iter()
351                    .map(|v| Cell::new(v).style_spec("lFc"))
352                    .collect::<Vec<Cell>>(),
353            ));
354        }
355        debug!("the remote list total: {}", self.0.len());
356        table.printstd();
357    }
358}
atsh_lib/connection/remote.rs

atsh_lib/connection/
remote.rs
