use crate::algorithm::hash;
use crate::algorithm::key_exchange::KeyExchange;
use crate::algorithm::public_key::PublicKey;
use crate::client::Client;
use crate::constant::{size, ssh_msg_code};
use crate::data::Data;
use crate::h::H;
use crate::packet::Packet;
use crate::window_size::WindowSize;
use crate::{kex, SshError, SshResult};
use std::io::{self, Read, Write};
impl<S> Client<S>
where
S: Read + Write,
{
pub(crate) fn read_version(&mut self) -> Vec<u8> {
let mut v = [0_u8; 128];
loop {
match self.stream.read(&mut v) {
Ok(i) => return v[..i].to_vec(),
Err(_) => continue,
};
}
}
pub fn read(&mut self) -> SshResult<Vec<Data>> {
self.read_data(None)
}
pub fn read_data(&mut self, lws: Option<&mut WindowSize>) -> SshResult<Vec<Data>> {
self.timeout.is_timeout()?;
let mut results = vec![];
let mut result = vec![0; size::BUF_SIZE as usize];
let len = match self.stream.read(&mut result) {
Ok(len) => {
if len == 0 {
return Ok(results);
}
self.timeout.renew();
len
}
Err(e) => {
if Client::<S>::is_would_block(&e) {
return Ok(results);
}
return Err(SshError::from(e));
}
};
result.truncate(len);
match self.is_encryption {
true => self.process_data_encrypt(result, &mut results, lws)?,
false => self.process_data(result, &mut results),
}
Ok(results)
}
pub fn process_data(&mut self, mut result: Vec<u8>, results: &mut Vec<Data>) {
self.sequence.server_auto_increment();
let packet_len = &result[..4];
let mut packet_len_slice = [0_u8; 4];
packet_len_slice.copy_from_slice(packet_len);
let packet_len = (u32::from_be_bytes(packet_len_slice) as usize) + 4;
if result.len() > packet_len {
let (v1, v2) = result.split_at_mut(packet_len);
let data = Packet::from(v1.to_vec()).unpacking();
results.push(data);
result = v2.to_vec();
}
let data = Packet::from(result).unpacking();
results.push(data);
}
pub fn process_data_encrypt(
&mut self,
mut result: Vec<u8>,
results: &mut Vec<Data>,
mut lws: Option<&mut WindowSize>,
) -> SshResult<()> {
loop {
self.sequence.server_auto_increment();
if result.len() < 4 {
self.check_result_len(&mut result)?;
}
let data_len = {
self.encryption
.as_mut()
.unwrap()
.data_len(self.sequence.server_sequence_num, result.as_slice())
};
if result.len() < data_len {
self.get_encrypt_data(&mut result, data_len)?;
}
let (this, remaining) = result.split_at_mut(data_len);
let decryption_result = {
self.encryption
.as_mut()
.unwrap()
.decrypt(self.sequence.server_sequence_num, &mut this.to_vec())
}?;
let data = Packet::from(decryption_result).unpacking();
if !self.window_adjust(data.clone(), &mut lws)?
&& self.r_size_one_gb(data.clone(), &mut lws)?
{
results.push(data);
}
if remaining.is_empty() {
break;
}
result = remaining.to_vec();
}
Ok(())
}
fn get_encrypt_data(&mut self, result: &mut Vec<u8>, data_len: usize) -> SshResult<()> {
loop {
self.timeout.is_timeout()?;
let mut buf = vec![0; data_len - result.len()];
match self.stream.read(&mut buf) {
Ok(len) => {
if len > 0 {
self.timeout.renew();
buf.truncate(len);
result.extend(buf);
}
if result.len() >= data_len {
return Ok(());
}
}
Err(e) => {
if e.kind() == io::ErrorKind::WouldBlock {
continue;
}
return Err(SshError::from(e));
}
};
}
}
fn check_result_len(&mut self, result: &mut Vec<u8>) -> SshResult<usize> {
loop {
self.timeout.is_timeout()?;
let mut buf = vec![0; size::BUF_SIZE as usize];
match self.stream.read(&mut buf) {
Ok(len) => {
if len > 0 {
self.timeout.renew();
buf.truncate(len);
result.extend(buf);
}
if result.len() >= 4 {
return Ok(len);
}
}
Err(e) => {
if e.kind() == io::ErrorKind::WouldBlock {
continue;
}
return Err(SshError::from(e));
}
};
}
}
fn window_adjust(
&mut self,
mut data: Data,
lws: &mut Option<&mut WindowSize>,
) -> SshResult<bool> {
if let Some(v) = lws {
v.process_local_window_size(data.as_slice(), self)?;
let mc = data[0];
if mc == ssh_msg_code::SSH_MSG_CHANNEL_WINDOW_ADJUST {
data.get_u8();
data.get_u32();
let size = data.get_u32();
v.add_remote_window_size(size);
return Ok(true);
}
}
Ok(false)
}
fn r_size_one_gb(
&mut self,
mut data: Data,
lws: &mut Option<&mut WindowSize>,
) -> SshResult<bool> {
if self.is_w_1_gb {
return Ok(true);
}
return match data[0] {
ssh_msg_code::SSH_MSG_KEXINIT => {
self.is_r_1_gb = true;
log::info!("start for key negotiation.");
let mut h = H::new();
let cv = self.config.version.client_version.as_str();
let sv = self.config.version.server_version.as_str();
h.set_v_c(cv);
h.set_v_s(sv);
h.set_i_s(data.clone().as_slice());
kex::processing_server_algorithm(&mut self.config, data)?;
match lws {
None => kex::send_algorithm(&mut h, self, None)?,
Some(ws) => kex::send_algorithm(&mut h, self, Some(ws))?,
}
let key_exchange = self.config.algorithm.matching_key_exchange_algorithm()?;
let public_key = self.config.algorithm.matching_public_key_algorithm()?;
match lws {
None => kex::send_qc(self, key_exchange.get_public_key(), None)?,
Some(ws) => kex::send_qc(self, key_exchange.get_public_key(), Some(ws))?,
}
self.signature = Some(Signature {
h,
key_exchange,
public_key,
});
Ok(false)
}
ssh_msg_code::SSH_MSG_KEXDH_REPLY => {
let k = self.signature.as_mut().unwrap();
data.get_u8();
let sig = kex::generate_signature(data, &mut k.h, &mut k.key_exchange)?;
let session_id = hash::digest(&k.h.as_bytes(), k.key_exchange.get_hash_type());
let flag = k.public_key.verify_signature(&k.h.k_s, &session_id, &sig)?;
if !flag {
log::error!("signature verification failure.");
return Err(SshError::from("signature verification failure."));
}
log::info!("signature verification success.");
Ok(false)
}
ssh_msg_code::SSH_MSG_NEWKEYS => {
match lws {
None => kex::new_keys(self, None)?,
Some(ws) => kex::new_keys(self, Some(ws))?,
}
let k = self.signature.as_mut().unwrap();
let hash_type = k.key_exchange.get_hash_type();
let hash = hash::hash::Hash::new(k.h.clone(), &self.session_id, hash_type);
let mac = self.config.algorithm.matching_mac_algorithm()?;
let encryption = self
.config
.algorithm
.matching_encryption_algorithm(hash, mac)?;
self.encryption = Some(encryption);
self.is_encryption = true;
self.signature = None;
self.is_r_1_gb = false;
log::info!("key negotiation successful.");
Ok(false)
}
_ => Ok(true),
};
}
}
pub(crate) struct Signature {
h: H,
key_exchange: Box<dyn KeyExchange>,
public_key: Box<dyn PublicKey>,
}