use std::{collections::HashMap, str::FromStr};
use maybe_async::maybe_async;
use crate::{
packets::{
dfsc::{ReferralEntry, ReferralEntryValue},
smb2::Status,
},
Connection, Error, FileCreateArgs, Resource, Session, Tree,
};
use super::{config::ClientConfig, unc_path::UncPath};
pub struct Client {
config: ClientConfig,
connections: HashMap<UncPath, OpenedConnectionInfo>,
}
struct OpenedConnectionInfo {
conn: Connection,
session: Session,
tree: Tree,
creds: Option<(String, String)>,
}
impl Client {
pub fn new(config: ClientConfig) -> Self {
Client {
config,
connections: HashMap::new(),
}
}
#[maybe_async]
pub async fn close(&mut self) -> crate::Result<()> {
self.connections.clear();
Ok(())
}
pub fn list_shares(&self, server: &str) -> crate::Result<Vec<String>> {
unimplemented!()
}
#[maybe_async]
pub async fn share_connect(
&mut self,
unc: &UncPath,
user_name: &str,
password: String,
) -> crate::Result<()> {
if unc.share.is_none() {
return Err(crate::Error::InvalidArgument(
"UNC path does not contain a share name.".to_string(),
));
}
let share_unc = unc.clone().with_no_path();
if let Some(_) = self.connections.get(&share_unc) {
log::warn!("Connection already exists for this UNC path. Reusing it.");
}
let mut conn = Connection::build(share_unc.server.clone(), self.config.connection.clone())?;
conn.connect().await?;
let session = conn.authenticate(user_name, password.clone()).await?;
let tree = session.tree_connect(&share_unc.to_string()).await?;
let mut opened_conn_info = OpenedConnectionInfo {
conn,
session,
tree,
creds: None,
};
if self.config.dfs {
opened_conn_info.creds = Some((user_name.to_string(), password));
}
log::debug!("Connected to share {} with user {}", share_unc, user_name);
self.connections.insert(share_unc, opened_conn_info);
Ok(())
}
fn get_opened_conn_for_path(&self, unc: &UncPath) -> crate::Result<&OpenedConnectionInfo> {
if let Some(cst) = self.connections.get(&unc.clone().with_no_path()) {
Ok(cst)
} else {
Err(crate::Error::InvalidArgument(format!(
"No connection found for {}. Use `share_connect` to create one.",
unc
)))
}
}
#[maybe_async]
async fn create_file_shallow(
&self,
path: &UncPath,
args: &FileCreateArgs,
) -> crate::Result<Resource> {
let conn_info = self.get_opened_conn_for_path(path)?;
conn_info
.tree
.create(path.path.as_ref().map(|x| x.as_str()).unwrap_or(""), args)
.await
}
#[maybe_async]
pub async fn create_file(
&mut self,
path: &UncPath,
args: &FileCreateArgs,
) -> crate::Result<Resource> {
let file_result = self.create_file_shallow(path, args).await;
let resource = match file_result {
Ok(file) => Ok(file),
Err(Error::ReceivedErrorMessage(Status::U32_PATH_NOT_COVERED, _)) => {
if self.config.dfs {
DfsResolver::new(self).create_dfs_file(path, args).await
} else {
Err(Error::UnsupportedOperation(
"DFS is not enabled, but the server returned path not covered (dfs must be enabled in config to resolve the path!).".to_string(),
))
}
}
x => x,
}?;
Ok(resource)
}
}
struct DfsResolver<'a>(&'a mut Client);
impl<'a> DfsResolver<'a> {
fn new(client: &'a mut Client) -> Self {
DfsResolver(client)
}
#[maybe_async]
async fn create_dfs_file(
&mut self,
dfs_path: &UncPath,
args: &FileCreateArgs,
) -> crate::Result<Resource> {
let dfs_ref_paths = self.get_dfs_refs(dfs_path).await?;
let dfs_creds = self
.0
.get_opened_conn_for_path(dfs_path)?
.creds
.clone()
.ok_or_else(|| {
Error::InvalidState(
"DFS referral requires credentials, but none were found.".to_string(),
)
})?;
for ref_unc_path in dfs_ref_paths.iter() {
match self
.0
.share_connect(ref_unc_path, dfs_creds.0.as_str(), dfs_creds.1.clone())
.await
{
Err(e) => {
log::error!("Failed to open DFS referral: {}", e);
continue;
}
_ => (),
};
let resource = self
.0
.create_file_shallow(ref_unc_path, args)
.await
.map_err(|e| {
log::error!("Failed to create file on DFS referral: {}", e);
e
})?;
log::info!(
"Successfully created file on DFS referral: {}",
ref_unc_path
);
return Ok(resource);
}
Err(Error::DfsReferralConnectionFail(dfs_path.clone()))
}
#[maybe_async]
async fn get_dfs_refs(&self, unc: &UncPath) -> crate::Result<Vec<UncPath>> {
log::debug!("Resolving DFS referral for {}", unc);
let dfs_path_string = unc.to_string();
let dfs_root = self.0.get_opened_conn_for_path(unc)?.tree.as_dfs_tree()?;
let dfs_refs = dfs_root.dfs_get_referrals(&dfs_path_string).await?;
if !dfs_refs.referral_header_flags.storage_servers() {
return Err(Error::InvalidMessage(
"DFS referral does not contain storage servers".to_string(),
)
.into());
}
let mut paths = vec![];
for (indx, curr_referral) in dfs_refs.referral_entries.iter().enumerate() {
let is_first = indx == 0;
paths.push(self.ref_entry_to_dfs_target(
curr_referral,
dfs_refs.path_consumed as usize,
&dfs_path_string,
is_first,
)?);
}
Ok(paths)
}
fn ref_entry_to_dfs_target(
&self,
entry: &ReferralEntry,
path_consumed: usize,
dfs_path_string: &String,
is_first: bool,
) -> crate::Result<UncPath> {
match &entry.value {
ReferralEntryValue::V4(v4) => {
if v4.referral_entry_flags == 0 && is_first {
return Err(Error::InvalidMessage(
"First DFS Referral is not primary one, invalid message!".to_string(),
)
.into());
}
let index_end_of_match = path_consumed / std::mem::size_of::<u16>();
if index_end_of_match > dfs_path_string.len() {
return Err(Error::InvalidMessage(
"DFS path consumed is out of bounds".to_string(),
)
.into());
}
let suffix = if index_end_of_match < dfs_path_string.len() {
dfs_path_string
.char_indices()
.nth(index_end_of_match)
.ok_or_else(|| {
Error::InvalidMessage("DFS path consumed is out of bounds".to_string())
})?
.0
} else {
dfs_path_string.len()
};
let unc_str_dest = "\\".to_string()
+ &v4.refs.network_address.to_string()
+ &dfs_path_string[suffix..];
let unc_path = UncPath::from_str(&unc_str_dest)?;
log::debug!("Resolved DFS referral to {}", unc_path.to_string());
Ok(unc_path)
}
_ => {
return Err(Error::UnsupportedOperation(
"Unsupported DFS referral entry type".to_string(),
)
.into());
}
}
}
}