forked from NationTech/harmony
152 lines
3.9 KiB
Rust
152 lines
3.9 KiB
Rust
use crate::error::Error;
|
|
use async_trait::async_trait;
|
|
use log::info;
|
|
use russh::{
|
|
client::{Config as SshConfig, Handler, Msg},
|
|
Channel,
|
|
};
|
|
use russh_keys::key::{self, KeyPair};
|
|
use std::{fs, net::Ipv4Addr, sync::Arc};
|
|
|
|
#[async_trait]
|
|
pub trait ConfigRepository: std::fmt::Debug {
|
|
async fn load(&self) -> Result<String, Error>;
|
|
async fn save(&self, content: &str) -> Result<(), Error>;
|
|
}
|
|
|
|
struct Client {}
|
|
|
|
#[async_trait]
|
|
impl Handler for Client {
|
|
type Error = Error;
|
|
|
|
async fn check_server_key(
|
|
&mut self,
|
|
_server_public_key: &key::PublicKey,
|
|
) -> Result<bool, Self::Error> {
|
|
Ok(true)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum SshCredentials {
|
|
SshKey { username: String, key: Arc<KeyPair> },
|
|
Password { username: String, password: String },
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct SshConfigRepository {
|
|
ssh_config: Arc<SshConfig>,
|
|
credentials: SshCredentials,
|
|
host: (Ipv4Addr, u16),
|
|
}
|
|
|
|
impl SshConfigRepository {
|
|
pub fn new(
|
|
host: (Ipv4Addr, u16),
|
|
credentials: SshCredentials,
|
|
ssh_config: Arc<SshConfig>,
|
|
) -> Self {
|
|
Self {
|
|
ssh_config,
|
|
credentials,
|
|
host,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl SshConfigRepository {
|
|
async fn get_ssh_channel(&self) -> Result<Channel<Msg>, Error> {
|
|
let mut ssh = russh::client::connect(self.ssh_config.clone(), self.host, Client {}).await?;
|
|
|
|
match &self.credentials {
|
|
SshCredentials::SshKey { username, key } => {
|
|
ssh.authenticate_publickey(username, key.clone()).await?;
|
|
}
|
|
SshCredentials::Password { username, password } => {
|
|
ssh.authenticate_password(username, password).await?;
|
|
}
|
|
}
|
|
|
|
Ok(ssh.channel_open_session().await?)
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ConfigRepository for SshConfigRepository {
|
|
async fn load(&self) -> Result<String, Error> {
|
|
let mut channel = self.get_ssh_channel().await?;
|
|
|
|
channel.exec(true, "cat /conf/config.xml").await?;
|
|
let mut output: Vec<u8> = vec![];
|
|
loop {
|
|
let Some(msg) = channel.wait().await else {
|
|
break;
|
|
};
|
|
|
|
info!("got msg {:?}", msg);
|
|
match msg {
|
|
russh::ChannelMsg::Data { ref data } => {
|
|
output.append(&mut data.to_vec());
|
|
}
|
|
russh::ChannelMsg::ExitStatus { .. } => {}
|
|
russh::ChannelMsg::WindowAdjusted { .. } => {}
|
|
russh::ChannelMsg::Success { .. } => {}
|
|
russh::ChannelMsg::Eof { .. } => {}
|
|
_ => todo!(),
|
|
}
|
|
}
|
|
Ok(String::from_utf8(output).expect("Valid utf-8 bytes"))
|
|
}
|
|
|
|
async fn save(&self, content: &str) -> Result<(), Error> {
|
|
todo!("Backup, Validate, Reload config file");
|
|
let mut channel = self.get_ssh_channel().await?;
|
|
|
|
let command = format!(
|
|
"echo '{}' > /conf/config.xml",
|
|
content.replace("'", "'\"'\"'")
|
|
);
|
|
channel.exec(true, command.as_bytes()).await?;
|
|
|
|
loop {
|
|
let Some(msg) = channel.wait().await else {
|
|
break;
|
|
};
|
|
|
|
match msg {
|
|
russh::ChannelMsg::ExitStatus { exit_status } => {
|
|
if exit_status != 0 {
|
|
return Err(Error::Ssh(russh::Error::Disconnect));
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct LocalFileConfigRepository {
|
|
file_path: String,
|
|
}
|
|
|
|
impl LocalFileConfigRepository {
|
|
pub fn new(file_path: String) -> Self {
|
|
Self { file_path }
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ConfigRepository for LocalFileConfigRepository {
|
|
async fn load(&self) -> Result<String, Error> {
|
|
Ok(fs::read_to_string(&self.file_path)?)
|
|
}
|
|
|
|
async fn save(&self, content: &str) -> Result<(), Error> {
|
|
Ok(fs::write(&self.file_path, content)?)
|
|
}
|
|
}
|