use super::crypt_translater::{Arc4, ChaCha, CryptTranslator};
use super::error::Error;
use crypto::digest::Digest;
use crypto::sha2::Sha256;
use hex;
use std::io::prelude::*;
use std::net::TcpStream;
pub struct WireChannel {
stream: TcpStream,
read_buf: Vec<u8>,
read_trans: Option<Box<dyn CryptTranslator>>,
write_trans: Option<Box<dyn CryptTranslator>>,
}
impl WireChannel {
pub fn new(host: &str, port: u16) -> Result<WireChannel, Error> {
let stream = TcpStream::connect(format!("{}:{}", host, port))?;
Ok(WireChannel {
stream,
read_buf: Vec::new(),
read_trans: None,
write_trans: None,
})
}
pub fn set_crypt_key(&mut self, plugin: &[u8], key: &[u8], nonce: &[u8]) {
if plugin == b"ChaCha64" || plugin == b"ChaCha" {
let mut hasher = Sha256::new();
hasher.input(&key);
let key = &hex::decode(hasher.result_str()).unwrap();
self.read_trans = Some(Box::new(ChaCha::new(key, nonce)));
self.write_trans = Some(Box::new(ChaCha::new(key, nonce)));
} else if plugin == b"Arc4" {
self.read_trans = Some(Box::new(Arc4::new(key)));
self.write_trans = Some(Box::new(Arc4::new(key)));
}
}
pub fn read(&mut self, n: usize) -> Result<Vec<u8>, Error> {
let mut input_buf = [0u8; 4096];
while self.read_buf.len() < n {
let ln = self.stream.read(&mut input_buf)?;
self.read_buf.extend(input_buf[..ln].iter().copied());
}
let mut v: Vec<u8> = Vec::new();
for _ in 0..n {
v.push(self.read_buf.remove(0));
}
if let Some(ref mut trans) = self.read_trans {
let translated: Vec<u8> = (*trans.translate(&v)).to_vec();
Ok(translated)
} else {
Ok(v)
}
}
pub fn write(&mut self, buf: &[u8]) -> Result<(), Error> {
if let Some(ref mut trans) = self.write_trans {
self.stream.write(&*trans.translate(buf))?;
} else {
self.stream.write(buf)?;
}
Ok(())
}
}