ssh_test_server/
session.rs

1use crate::{command, SshExecuteHandler, UsersMap};
2use anyhow::Result;
3use async_trait::async_trait;
4use russh::server::{Auth, Handler, Msg, Response, Session};
5use russh::{Channel, ChannelId, ChannelMsg, CryptoVec};
6use russh_keys::key::PublicKey;
7use std::collections::HashMap;
8use std::mem;
9use std::sync::Arc;
10use tracing::debug;
11
12pub type ProgramsMap = Arc<HashMap<String, Box<SshExecuteHandler>>>;
13
14pub(crate) struct SshConnection {
15    id: usize,
16    users: UsersMap,
17    user: Option<String>,
18    programs: ProgramsMap,
19}
20
21impl SshConnection {
22    pub fn new(id: usize, users: UsersMap, programs: ProgramsMap) -> Self {
23        Self {
24            id,
25            users,
26            user: None,
27            programs,
28        }
29    }
30}
31
32#[async_trait]
33impl Handler for SshConnection {
34    type Error = anyhow::Error;
35
36    async fn auth_none(&mut self, user: &str) -> Result<Auth, Self::Error> {
37        debug!("auth_none user={user}");
38        Ok(Auth::Reject {
39            proceed_with_methods: None,
40        })
41    }
42
43    async fn auth_password(&mut self, user: &str, password: &str) -> Result<Auth, Self::Error> {
44        let users = self.users.lock().unwrap();
45        if let Some(u) = users.get(user) {
46            if password == u.password() {
47                self.user = Some(u.login().to_string());
48                drop(users);
49                debug!("auth_password user={user} password={password} Accepted");
50                return Ok(Auth::Accept);
51            }
52        }
53
54        drop(users);
55        debug!("auth_password user={user} password={password} Rejected");
56        Ok(Auth::Reject {
57            proceed_with_methods: None,
58        })
59    }
60
61    async fn auth_publickey(
62        &mut self,
63        user: &str,
64        public_key: &PublicKey,
65    ) -> Result<Auth, Self::Error> {
66        debug!("auth_publickey user={user} public_key={public_key:?}");
67
68        Ok(Auth::Reject {
69            proceed_with_methods: None,
70        })
71    }
72
73    async fn auth_keyboard_interactive(
74        &mut self,
75        user: &str,
76        submethods: &str,
77        _response: Option<Response<'async_trait>>,
78    ) -> Result<Auth, Self::Error> {
79        debug!("auth_keyboard_interactive user={user} submethods={submethods:?}");
80        Ok(Auth::Reject {
81            proceed_with_methods: None,
82        })
83    }
84
85    async fn auth_succeeded(&mut self, _session: &mut Session) -> Result<(), Self::Error> {
86        debug!("auth_succeeded");
87        Ok(())
88    }
89
90    async fn channel_close(
91        &mut self,
92        channel: ChannelId,
93        _session: &mut Session,
94    ) -> Result<(), Self::Error> {
95        debug!("channel_close channel={channel}");
96        Ok(())
97    }
98
99    async fn channel_open_session(
100        &mut self,
101        mut channel: Channel<Msg>,
102        session: &mut Session,
103    ) -> Result<bool, Self::Error> {
104        let session_id = self.id;
105        debug!(session_id, "channel_open_session channel={}", channel.id());
106        let handle = session.handle();
107        let user = self.user.clone().unwrap();
108        let users = self.users.clone();
109        let programs = self.programs.clone();
110        tokio::spawn(async move {
111            let id = channel.id();
112            let mut command_buf = vec![];
113
114            while let Some(msg) = channel.wait().await {
115                match msg {
116                    ChannelMsg::RequestPty {
117                        want_reply,
118                        term,
119                        col_width,
120                        row_height,
121                        pix_width,
122                        pix_height,
123                        terminal_modes,
124                    } => {
125                        debug!(session_id, "request-pty want_reply={want_reply} term={term} col/row={col_width}/{row_height} pix width/height={pix_width}/{pix_height} modes={terminal_modes:?}");
126                        if want_reply {
127                            handle.channel_success(id).await.unwrap();
128                        }
129                    }
130                    ChannelMsg::RequestShell { want_reply } => {
131                        debug!(session_id, "request-shell want_reply={want_reply}");
132                        if want_reply {
133                            handle.channel_success(id).await.unwrap();
134                        }
135                        handle.data(id, CryptoVec::from_slice(b"$ ")).await.unwrap();
136                    }
137                    ChannelMsg::Data { data } => {
138                        debug!(session_id, "data={}", String::from_utf8_lossy(&data));
139
140                        let mut stdout = CryptoVec::new();
141                        for b in data.iter() {
142                            if *b == 0x03 {
143                                // Ctrl + C
144                                handle.exit_status_request(id, 130).await.unwrap();
145                                handle.close(id).await.unwrap();
146                            } else if *b == b'\r' || *b == b'\n' {
147                                stdout.push(b'\r');
148                                stdout.push(b'\n');
149                                handle.data(id, mem::take(&mut stdout)).await.unwrap();
150                                let cmd = mem::take(&mut command_buf);
151                                command::execute_command(
152                                    cmd, id, &handle, &user, &users, &programs,
153                                )
154                                .await;
155                                handle.data(id, CryptoVec::from_slice(b"$ ")).await.unwrap();
156                            } else {
157                                command_buf.push(*b);
158                                stdout.push(*b);
159                            }
160                        }
161
162                        if !stdout.is_empty() {
163                            handle.data(id, mem::take(&mut stdout)).await.unwrap();
164                        }
165                    }
166                    ChannelMsg::Exec {
167                        want_reply,
168                        command,
169                    } => {
170                        debug!(
171                            session_id,
172                            "exec want_reply={want_reply} command: {}",
173                            String::from_utf8_lossy(&command)
174                        );
175                        if want_reply {
176                            handle.channel_success(id).await.unwrap();
177                        }
178
179                        command::execute_command(command, id, &handle, &user, &users, &programs)
180                            .await;
181                        handle.close(id).await.unwrap();
182                    }
183                    _ => {
184                        debug!(session_id, "msg={msg:?}");
185                    }
186                }
187            }
188            debug!(session_id, "closed");
189        });
190
191        Ok(true)
192    }
193
194    async fn channel_open_x11(
195        &mut self,
196        channel: Channel<Msg>,
197        originator_address: &str,
198        originator_port: u32,
199        _session: &mut Session,
200    ) -> Result<bool, Self::Error> {
201        debug!("channel_open_x11 channel={} originator_address={originator_address} originator_port={originator_port}", channel.id());
202        Ok(false)
203    }
204
205    async fn channel_open_direct_tcpip(
206        &mut self,
207        channel: Channel<Msg>,
208        host_to_connect: &str,
209        port_to_connect: u32,
210        originator_address: &str,
211        originator_port: u32,
212        _session: &mut Session,
213    ) -> Result<bool, Self::Error> {
214        debug!("channel_open_direct_tcpip channel={} host_to_connect={host_to_connect} port_to_connect={port_to_connect} originator_address={originator_address} originator_port={originator_port}", channel.id());
215        Ok(false)
216    }
217
218    async fn channel_open_forwarded_tcpip(
219        &mut self,
220        channel: Channel<Msg>,
221        host_to_connect: &str,
222        port_to_connect: u32,
223        originator_address: &str,
224        originator_port: u32,
225        _session: &mut Session,
226    ) -> Result<bool, Self::Error> {
227        debug!("channel_open_forwarded_tcpip channel={} host_to_connect={host_to_connect} port_to_connect={port_to_connect} originator_address={originator_address} originator_port={originator_port}", channel.id());
228        Ok(false)
229    }
230
231    fn adjust_window(&mut self, channel: ChannelId, current: u32) -> u32 {
232        debug!("adjust_window {channel} current={current}");
233        current
234    }
235}