mini_telnet/
lib.rs

1mod codec;
2pub mod error;
3
4use encoding::DecoderTrap;
5use encoding::{all::GB18030, all::GBK, Encoding};
6use futures::stream::StreamExt;
7use regex::bytes::Regex;
8use tokio::{
9    io::AsyncWriteExt,
10    net::TcpStream,
11    time::{self, Duration},
12};
13use tokio_util::codec::FramedRead;
14
15use crate::codec::{Item, TelnetCodec};
16use crate::error::TelnetError;
17
18#[derive(Debug, Default)]
19pub struct TelnetBuilder {
20    prompts: Vec<String>,
21    username_prompt: String,
22    password_prompt: String,
23    connect_timeout: Duration,
24    timeout: Duration,
25}
26
27impl TelnetBuilder {
28    /// Set the telnet server prompt, as many characters as possible.(`~` or `#` is not good. May misjudge).
29    pub fn prompt<T: ToString>(mut self, prompt: T) -> TelnetBuilder {
30        self.prompts = vec![prompt.to_string()];
31        self
32    }
33
34    /// Set the telnet server prompts, as many characters as possible.(`~` or `#` is not good. May misjudge).
35    /// If `prompts` is set, `prompt` will be overwritten.
36    pub fn prompts<T: ToString>(mut self, prompts: &[T]) -> TelnetBuilder {
37        self.prompts = prompts.iter().map(|p| p.to_string()).collect();
38        self
39    }
40
41    /// Login prompt, the common ones are `login: ` and `Password: ` or `Username:` and `Password:`.
42    pub fn login_prompt(mut self, user_prompt: &str, pass_prompt: &str) -> TelnetBuilder {
43        self.username_prompt = user_prompt.to_string();
44        self.password_prompt = pass_prompt.to_string();
45        self
46    }
47
48    /// Set the timeout for `TcpStream` connect remote addr.
49    pub fn connect_timeout(mut self, connect_timeout: Duration) -> TelnetBuilder {
50        self.connect_timeout = connect_timeout;
51        self
52    }
53
54    /// Set the timeout for the operation.
55    pub fn timeout(mut self, timeout: Duration) -> TelnetBuilder {
56        self.timeout = timeout;
57        self
58    }
59
60    /// Establish a connection with the remote telnetd.
61    pub async fn connect(self, addr: &str) -> Result<Telnet, TelnetError> {
62        let clear = Clear::new()?;
63        match time::timeout(self.connect_timeout, TcpStream::connect(addr)).await {
64            Ok(res) => Ok(Telnet {
65                content: vec![],
66                stream: res?,
67                timeout: self.timeout,
68                prompts: self.prompts,
69                username_prompt: self.username_prompt,
70                password_prompt: self.password_prompt,
71                clear,
72            }),
73            Err(_) => Err(TelnetError::Timeout(format!(
74                "Connect remote addr({})",
75                addr
76            ))),
77        }
78    }
79}
80
81pub struct Telnet {
82    timeout: Duration,
83    content: Vec<String>,
84    stream: TcpStream,
85    prompts: Vec<String>,
86    username_prompt: String,
87    password_prompt: String,
88    clear: Clear,
89}
90
91impl Telnet {
92    /// Create a `TelnetBuilder`
93    pub fn builder() -> TelnetBuilder {
94        TelnetBuilder::default()
95    }
96    // Format the end of the string as a `\n`
97    fn format_enter_str(s: &str) -> String {
98        if !s.ends_with('\n') {
99            format!("{}\n", s)
100        } else {
101            s.to_string()
102        }
103    }
104
105    /// Login remote telnet daemon, only retry one time.
106    /// # Examples
107    ///
108    /// ```no_run
109    /// let mut client = Telnet::builder()
110    ///     .prompt("username@hostname:$ ")
111    ///     .login_prompt("login: ", "Password: ")
112    ///     .connect_timeout(Duration::from_secs(3))
113    ///     .connect("192.168.0.1:23").await?;
114    ///
115    /// match client.login("username", "password").await {
116    ///     Ok(_) => println!("login success."),
117    ///     Err(e) => println!("login failed: {}", e),
118    /// };
119    /// ```
120    ///
121    pub async fn login(&mut self, username: &str, password: &str) -> Result<(), TelnetError> {
122        let user = Telnet::format_enter_str(username);
123        let pass = Telnet::format_enter_str(password);
124
125        // Only retry one time, if password is input, then set with `true`;
126        let mut auth_failed = false;
127
128        let (read, mut write) = self.stream.split();
129        let mut telnet = FramedRead::new(read, TelnetCodec::default());
130
131        loop {
132            match time::timeout(self.timeout, telnet.next()).await {
133                Ok(res) => {
134                    match res {
135                        Some(res) => {
136                            match res? {
137                                Item::Do(i) | Item::Dont(i) => {
138                                    // set window size
139                                    if i == 0x1f {
140                                        write
141                                            .write_all(&[
142                                                0xff, 0xfb, 0x1f, 0xff, 0xfa, 0x1f, 0x00, 0xfc,
143                                                0x00, 0x1b, 0xff, 0xf0,
144                                            ])
145                                            .await?;
146                                    } else {
147                                        write.write_all(&[0xff, 0xfc, i]).await?;
148                                    }
149                                }
150                                Item::Will(i) | Item::Wont(i) => {
151                                    write.write_all(&[0xff, 0xfe, i]).await?;
152                                }
153                                Item::Line(line) => {
154                                    let line = self.clear.color(&line);
155                                    if line.ends_with(self.username_prompt.as_bytes()) {
156                                        if auth_failed {
157                                            return Err(TelnetError::AuthenticationFailed);
158                                        }
159                                        write.write_all(user.as_bytes()).await?;
160                                    } else if line.ends_with(self.password_prompt.as_bytes()) {
161                                        write.write_all(pass.as_bytes()).await?;
162                                        auth_failed = true;
163                                    } else if self
164                                        .prompts
165                                        .iter()
166                                        .filter(|p| line.ends_with(p.as_bytes()))
167                                        .count()
168                                        != 0
169                                    {
170                                        return Ok(());
171                                    }
172                                }
173                                item => return Err(TelnetError::UnknownIAC(format!("{:?}", item))),
174                            }
175                        }
176                        None => return Err(TelnetError::NoMoreData),
177                    };
178                }
179                Err(_) => return Err(TelnetError::Timeout("login".to_string())),
180            }
181        }
182    }
183
184    /// Execute command, and filter it input message by line count.
185    ///
186    /// # Examples
187    ///
188    /// ```no_run
189    ///assert_eq!(telnet.execute("echo 'haha'").await?, "haha\n");
190    /// ```
191    ///
192    pub async fn execute(&mut self, cmd: &str) -> Result<String, TelnetError> {
193        let command = Telnet::format_enter_str(cmd);
194        let mut incomplete_line: Vec<u8> = vec![];
195        let mut line_feed_cnt = command.lines().count() as isize;
196        let mut real_output = false;
197
198        let (read, mut write) = self.stream.split();
199        match time::timeout(self.timeout, write.write(command.as_bytes())).await {
200            Ok(res) => res?,
201            Err(_) => return Err(TelnetError::Timeout("write cmd".to_string())),
202        };
203        let mut telnet = FramedRead::new(read, TelnetCodec::default());
204
205        loop {
206            match time::timeout(self.timeout, telnet.next()).await {
207                Ok(res) => match res {
208                    Some(item) => {
209                        if let Item::Line(line) = item? {
210                            let mut line = self.clear.color(&line);
211
212                            // ignore prompt line
213                            if self
214                                .prompts
215                                .iter()
216                                .filter(|p| line.ends_with(p.as_bytes()))
217                                .count()
218                                != 0
219                            {
220                                break;
221                            }
222                            // ignore command line echo
223                            if line.ends_with(&[10]) && line_feed_cnt > 0 {
224                                line_feed_cnt -= 1;
225                                if line_feed_cnt == 0 {
226                                    real_output = true;
227                                    continue;
228                                }
229                            }
230
231                            if !real_output {
232                                continue;
233                            }
234
235                            if !line.ends_with(&[10]) || !incomplete_line.is_empty() {
236                                incomplete_line.append(&mut line);
237                            } else {
238                                self.content.push(decode(&line)?);
239                                continue;
240                            }
241                            // ignore command line
242                            if self
243                                .prompts
244                                .iter()
245                                .filter(|p| incomplete_line.ends_with(p.as_bytes()))
246                                .count()
247                                != 0
248                            {
249                                break;
250                            }
251                            if incomplete_line.ends_with(&[10]) {
252                                self.content.push(decode(&incomplete_line)?);
253                                incomplete_line.clear();
254                            }
255                        }
256                    }
257                    None => return Err(TelnetError::NoMoreData),
258                },
259                Err(_) => return Err(TelnetError::Timeout("read next framed".to_string())),
260            }
261        }
262        let result = self.content.join("");
263        self.content.clear();
264        Ok(result)
265    }
266
267    /// All echoed content is returned when the command is executed.(**Note** that this may contain some
268    /// useless information, such as prompts, which need to be filtered and processed by yourself.)
269    ///
270    /// # Examples
271    ///
272    /// ```no_run
273    /// assert_eq!(
274    ///     "echo 'haha'\nhaha\n",
275    ///     telnet.normal_execute("echo 'haha'").await?
276    /// );
277    ///```
278    ///
279    pub async fn normal_execute(&mut self, cmd: &str) -> Result<String, TelnetError> {
280        let command = Telnet::format_enter_str(cmd);
281        let mut incomplete_line: Vec<u8> = vec![];
282
283        let (read, mut write) = self.stream.split();
284        match time::timeout(self.timeout, write.write(command.as_bytes())).await {
285            Ok(res) => res?,
286            Err(_) => return Err(TelnetError::Timeout("write cmd".to_string())),
287        };
288        let mut telnet = FramedRead::new(read, TelnetCodec::default());
289
290        loop {
291            match time::timeout(self.timeout, telnet.next()).await {
292                Ok(res) => match res {
293                    Some(item) => {
294                        if let Item::Line(line) = item? {
295                            let mut line = self.clear.color(&line);
296                            if self
297                                .prompts
298                                .iter()
299                                .filter(|p| line.ends_with(p.as_bytes()))
300                                .count()
301                                != 0
302                            {
303                                break;
304                            }
305
306                            if !line.ends_with(&[10]) || !incomplete_line.is_empty() {
307                                incomplete_line.append(&mut line);
308                            } else {
309                                self.content.push(decode(&line)?);
310                                continue;
311                            }
312                            // ignore command line
313                            if self
314                                .prompts
315                                .iter()
316                                .filter(|p| incomplete_line.ends_with(p.as_bytes()))
317                                .count()
318                                != 0
319                            {
320                                break;
321                            }
322                            if incomplete_line.ends_with(&[10]) {
323                                self.content.push(decode(&incomplete_line)?);
324                                incomplete_line.clear();
325                            }
326                        }
327                    }
328                    None => return Err(TelnetError::NoMoreData),
329                },
330                Err(_) => return Err(TelnetError::Timeout("read next framed".to_string())),
331            }
332        }
333        let result = self.content.join("");
334        self.content.clear();
335        Ok(result)
336    }
337}
338
339fn decode(line: &[u8]) -> Result<String, TelnetError> {
340    match String::from_utf8(line.to_vec()) {
341        Ok(result) => Ok(result),
342        Err(e) => {
343            if let Ok(result) = GBK.decode(line, DecoderTrap::Strict) {
344                return Ok(result);
345            }
346
347            if let Ok(result) = GB18030.decode(line, DecoderTrap::Strict) {
348                return Ok(result);
349            }
350            Err(TelnetError::ParseError(e))
351        }
352    }
353}
354
355struct Clear {
356    color_re: Regex,
357}
358
359impl Clear {
360    pub fn new() -> Result<Self, TelnetError> {
361        let color_re = Regex::new(r"\[\d{2,3}m")?;
362        Ok(Self { color_re })
363    }
364
365    pub fn color(&self, content: &[u8]) -> Vec<u8> {
366        self.color_re
367            .replace_all(content, &[] as &[u8])
368            .into_owned()
369    }
370}
mini_telnet/lib.rs

mini_telnet/
lib.rs
