use std::sync::Arc;
use crate::sync_helpers::*;
use maybe_async::*;
use crate::connection::negotiation_state::NegotiateState;
use crate::connection::preauth_hash::PreauthHashValue;
use crate::crypto::{
kbkdf_hmacsha256, make_encrypting_algo, make_signing_algo, CryptoError, DerivedKey, KeyToDerive,
};
use crate::packets::smb2::{EncryptionCipher, SessionFlags, SigningAlgorithmId};
use super::{MessageDecryptor, MessageEncryptor, MessageSigner};
#[derive(Debug)]
pub struct SessionState {
pub session_id: u64,
flags: SessionFlags,
signer: Option<MessageSigner>,
decryptor: Option<MessageDecryptor>,
encryptor: Option<MessageEncryptor>,
}
impl SessionState {
const SIGNING_KEY_LABEL: &[u8] = b"SMBSigningKey\x00";
const S2C_DECRYPTION_KEY_LABEL: &[u8] = b"SMBS2CCipherKey\x00";
const C2S_ENCRYPTION_KEY_LABEL: &[u8] = b"SMBC2SCipherKey\x00";
#[maybe_async]
pub async fn set(
state: &mut Arc<Mutex<Self>>,
session_key: &KeyToDerive,
preauth_hash: &PreauthHashValue,
negotation_state: &NegotiateState,
) -> crate::Result<()> {
let deriver = KeyDeriver::new(session_key, preauth_hash);
let signer = Self::make_signer(&deriver, negotation_state.signing_algo())?;
let decryptor = Self::make_decryptor(&deriver, negotation_state.cipher())?;
let encryptor = Self::make_encryptor(&deriver, negotation_state.cipher())?;
let mut state = state.lock().await?;
state.signer = Some(signer);
state.decryptor = Some(decryptor);
state.encryptor = Some(encryptor);
Ok(())
}
fn make_signer(
deriver: &KeyDeriver,
signing_algo: SigningAlgorithmId,
) -> Result<MessageSigner, CryptoError> {
let signing_key = deriver.derive(Self::SIGNING_KEY_LABEL)?;
Ok(MessageSigner::new(make_signing_algo(
signing_algo,
&signing_key,
)?))
}
fn make_encryptor(
deriver: &KeyDeriver,
cipher: EncryptionCipher,
) -> Result<MessageEncryptor, CryptoError> {
let c2s_encryption_key = deriver.derive(Self::C2S_ENCRYPTION_KEY_LABEL)?;
Ok(MessageEncryptor::new(make_encrypting_algo(
cipher,
&c2s_encryption_key,
)?))
}
fn make_decryptor(
deriver: &KeyDeriver,
cipher: EncryptionCipher,
) -> Result<MessageDecryptor, CryptoError> {
let s2c_decryption_key = deriver.derive(Self::S2C_DECRYPTION_KEY_LABEL)?;
Ok(MessageDecryptor::new(make_encrypting_algo(
cipher,
&s2c_decryption_key,
)?))
}
#[maybe_async]
pub async fn set_flags(state: &mut Arc<Mutex<Self>>, flags: SessionFlags) -> crate::Result<()> {
state.lock().await?.flags = flags;
Ok(())
}
#[maybe_async]
pub async fn invalidate(state: &Arc<Mutex<Self>>) -> crate::Result<()> {
let mut state = state.lock().await?;
state.signer = None;
state.decryptor = None;
state.encryptor = None;
Ok(())
}
#[maybe_async]
pub async fn signing_enabled(state: &Arc<Mutex<Self>>) -> crate::Result<bool> {
Ok(state.lock().await?.signer.is_some())
}
#[maybe_async]
pub async fn encryption_enabled(state: &Arc<Mutex<Self>>) -> crate::Result<bool> {
let state = state.lock().await?;
return Ok(state.flags.encrypt_data()
&& state.encryptor.is_some()
&& state.decryptor.is_some());
}
#[maybe_async]
pub async fn is_set_up(state: &Arc<Mutex<Self>>) -> crate::Result<bool> {
Ok(Self::encryption_enabled(state).await? || Self::signing_enabled(state).await?)
}
pub fn decryptor(&mut self) -> Option<&mut MessageDecryptor> {
self.decryptor.as_mut()
}
pub fn encryptor(&mut self) -> Option<&mut MessageEncryptor> {
self.encryptor.as_mut()
}
pub fn signer(&mut self) -> Option<&mut MessageSigner> {
self.signer.as_mut()
}
}
struct KeyDeriver<'a> {
session_key: &'a KeyToDerive,
preauth_hash: &'a PreauthHashValue,
}
impl<'a> KeyDeriver<'a> {
#[inline]
pub fn new(session_key: &'a KeyToDerive, preauth_hash: &'a PreauthHashValue) -> Self {
Self {
session_key,
preauth_hash,
}
}
#[inline]
pub fn derive(&self, label: &[u8]) -> Result<DerivedKey, CryptoError> {
kbkdf_hmacsha256::<16>(self.session_key, label, self.preauth_hash)
}
}
impl Default for SessionState {
fn default() -> Self {
Self {
session_id: Default::default(),
flags: SessionFlags::new(),
signer: Default::default(),
decryptor: Default::default(),
encryptor: Default::default(),
}
}
}