use std::io::{Read, Write};
use crate::algorithm::hash;
use crate::algorithm::key_exchange::KeyExchange;
use crate::algorithm::public_key::PublicKey;
use crate::config::{
CompressionAlgorithm, Config, EncryptionAlgorithm, KeyExchangeAlgorithm, MacAlgorithm,
PublicKeyAlgorithm,
};
use crate::constant::ssh_msg_code;
use crate::data::Data;
use crate::error::{SshError, SshResult};
use crate::slog::log;
use crate::window_size::WindowSize;
use crate::{client::Client, h::H, util};
pub(crate) fn send_algorithm<S>(
h: &mut H,
client: &mut Client<S>,
rws: Option<&mut WindowSize>,
) -> SshResult<()>
where
S: Read + Write,
{
log::info!(
"client algorithms: [{}]",
client.config.algorithm.client_algorithm.to_string()
);
let mut data = Data::new();
data.put_u8(ssh_msg_code::SSH_MSG_KEXINIT);
data.extend(util::cookie());
data.extend(client.config.algorithm.client_algorithm.as_i());
data.put_str("")
.put_str("")
.put_u8(false as u8)
.put_u32(0_u32);
h.set_i_c(data.clone().as_slice());
match rws {
None => client.write(data)?,
Some(ws) => client.write_data(data, Some(ws))?,
}
Ok(())
}
pub(crate) fn receive_algorithm<S>(
h: &mut H,
client: &mut Client<S>,
mut rws: Option<&mut WindowSize>,
) -> SshResult<()>
where
S: Read + Write,
{
loop {
let results = match &mut rws {
None => client.read()?,
Some(ws) => client.read_data(Some(ws))?,
};
for result in results {
if result.is_empty() {
continue;
}
let message_code = result[0];
if message_code == ssh_msg_code::SSH_MSG_KEXINIT {
h.set_i_s(result.as_slice());
processing_server_algorithm(&mut client.config, result)?;
return Ok(());
}
}
}
}
pub(crate) fn processing_server_algorithm(config: &mut Config, mut data: Data) -> SshResult<()> {
data.get_u8();
data.skip(16);
let server_algorithm = &mut config.algorithm.server_algorithm;
server_algorithm.key_exchange_algorithm =
KeyExchangeAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
server_algorithm.public_key_algorithm =
PublicKeyAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
server_algorithm.c_encryption_algorithm =
EncryptionAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
server_algorithm.s_encryption_algorithm =
EncryptionAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
server_algorithm.c_mac_algorithm = MacAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
server_algorithm.s_mac_algorithm = MacAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
server_algorithm.c_compression_algorithm =
CompressionAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
server_algorithm.s_compression_algorithm =
CompressionAlgorithm(util::vec_u8_to_string(data.get_u8s(), ",")?);
log::info!("server algorithms: [{}]", server_algorithm.to_string());
Ok(())
}
pub(crate) fn send_qc<S>(
client: &mut Client<S>,
public_key: &[u8],
rws: Option<&mut WindowSize>,
) -> SshResult<()>
where
S: Read + Write,
{
let mut data = Data::new();
data.put_u8(ssh_msg_code::SSH_MSG_KEXDH_INIT);
data.put_u8s(public_key);
match rws {
None => client.write(data),
Some(ws) => client.write_data(data, Some(ws)),
}
}
pub(crate) fn verify_signature_and_new_keys<S>(
client: &mut Client<S>,
public_key: &mut Box<dyn PublicKey>,
key_exchange: &mut Box<dyn KeyExchange>,
h: &mut H,
mut rws: Option<&mut WindowSize>,
) -> SshResult<Vec<u8>>
where
S: Read + Write,
{
let mut session_id = vec![];
loop {
let results = match &mut rws {
None => client.read()?,
Some(ws) => client.read_data(Some(ws))?,
};
for mut result in results {
if result.is_empty() {
continue;
}
let message_code = result.get_u8();
match message_code {
ssh_msg_code::SSH_MSG_KEXDH_REPLY => {
let sig = generate_signature(result, h, key_exchange)?;
session_id = hash::digest(&h.as_bytes(), key_exchange.get_hash_type());
let flag = public_key.verify_signature(&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.");
}
ssh_msg_code::SSH_MSG_NEWKEYS => {
match &mut rws {
None => new_keys(client, None)?,
Some(ws) => new_keys(client, Some(ws))?,
};
return Ok(session_id);
}
_ => {}
}
}
}
}
pub(crate) fn generate_signature(
mut data: Data,
h: &mut H,
key_exchange: &mut Box<dyn KeyExchange>,
) -> SshResult<Vec<u8>> {
let ks = data.get_u8s();
h.set_k_s(&ks);
let qs = data.get_u8s();
h.set_q_c(key_exchange.get_public_key());
h.set_q_s(&qs);
let vec = key_exchange.get_shared_secret(qs)?;
h.set_k(&vec);
let h = data.get_u8s();
let mut hd = Data::from(h);
hd.get_u8s();
let signature = hd.get_u8s();
Ok(signature)
}
pub(crate) fn new_keys<S>(
client: &mut Client<S>,
rws: Option<&mut WindowSize>,
) -> Result<(), SshError>
where
S: Read + Write,
{
let mut data = Data::new();
data.put_u8(ssh_msg_code::SSH_MSG_NEWKEYS);
match rws {
None => client.write(data)?,
Some(ws) => client.write_data(data, Some(ws))?,
}
log::info!("send new keys");
Ok(())
}
pub(crate) fn key_agreement<S>(
h: &mut H,
client: &mut Client<S>,
mut rws: Option<&mut WindowSize>,
) -> SshResult<hash::HashType>
where
S: Read + Write,
{
log::info!("start for key negotiation.");
log::info!("send client algorithm list.");
match &mut rws {
None => send_algorithm(h, client, None)?,
Some(ws) => send_algorithm(h, client, Some(ws))?,
}
log::info!("receive server algorithm list.");
match &mut rws {
None => receive_algorithm(h, client, None)?,
Some(ws) => receive_algorithm(h, client, Some(ws))?,
}
let mut key_exchange = client.config.algorithm.matching_key_exchange_algorithm()?;
let mut public_key = client.config.algorithm.matching_public_key_algorithm()?;
match &mut rws {
None => send_qc(client, key_exchange.get_public_key(), None)?,
Some(ws) => send_qc(client, key_exchange.get_public_key(), Some(ws))?,
}
let session_id = match &mut rws {
None => verify_signature_and_new_keys(client, &mut public_key, &mut key_exchange, h, None)?,
Some(ws) => {
verify_signature_and_new_keys(client, &mut public_key, &mut key_exchange, h, Some(ws))?
}
};
if client.session_id.is_empty() {
if session_id.is_empty() {
return Err(SshError::from("session id is none."));
}
client.session_id = session_id;
}
let hash_type = key_exchange.get_hash_type();
let hash = hash::hash::Hash::new(h.clone(), &client.session_id, hash_type);
let mac = client.config.algorithm.matching_mac_algorithm()?;
let encryption = client
.config
.algorithm
.matching_encryption_algorithm(hash, mac)?;
client.encryption = Some(encryption);
client.is_encryption = true;
log::info!("key negotiation successful.");
Ok(hash_type)
}