Skip to content

Commit

Permalink
Small improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
Kraemii committed Nov 13, 2024
1 parent 3cd85b2 commit 423af69
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 70 deletions.
6 changes: 6 additions & 0 deletions rust/src/nasl/builtin/network/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,12 @@ pub fn get_kb_item(context: &Context, name: &str) -> Result<Option<NaslValue>, F
.map_err(|e| e.into())
}

pub fn get_kb_item_str(context: &Context, name: &str) -> Result<String, FunctionErrorKind> {
get_kb_item(context, name)?
.map(|x| x.to_string())
.ok_or_else(|| FunctionErrorKind::Diagnostic(format!("KB key {} is not set", name), None))
}

pub fn verify_port(port: i64) -> Result<u16, FunctionErrorKind> {
if !(0..=65535).contains(&port) {
return Err(FunctionErrorKind::WrongArgument(format!(
Expand Down
75 changes: 22 additions & 53 deletions rust/src/nasl/builtin/network/socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use dns_lookup::lookup_host;
use nasl_function_proc_macro::nasl_function;

use super::{
get_kb_item, get_retry,
get_kb_item, get_kb_item_str, get_retry,
network_utils::{convert_timeout, ipstr2ipaddr},
tcp::TcpConnection,
tls::create_tls_client,
Expand Down Expand Up @@ -52,7 +52,7 @@ enum NaslSocket {
// This way the size of the enum is reduced
Tcp(Box<TcpConnection>),
Udp(UdpConnection),
Close,
Closed,
}

#[derive(Default)]
Expand Down Expand Up @@ -87,14 +87,14 @@ impl NaslSockets {
fn close(&self, socket_fd: usize) -> Result<NaslValue, FunctionErrorKind> {
let mut handles = self.handles.write().unwrap();
match handles.handles.get_mut(socket_fd) {
Some(NaslSocket::Close) => {
Some(NaslSocket::Closed) => {
return Err(FunctionErrorKind::Diagnostic(
"the given socket FD is already closed".to_string(),
None,
))
}
Some(socket) => {
*socket = NaslSocket::Close;
*socket = NaslSocket::Closed;
handles.closed_fd.push(socket_fd);
}
None => {
Expand Down Expand Up @@ -168,7 +168,7 @@ impl NaslSockets {
}
Ok(conn.write(data)?)
}
NaslSocket::Close => Err(FunctionErrorKind::WrongArgument(
NaslSocket::Closed => Err(FunctionErrorKind::WrongArgument(
"the given socket FD is already closed".to_string(),
)),
}
Expand Down Expand Up @@ -221,7 +221,7 @@ impl NaslSockets {

Ok(NaslValue::Data(data[..pos].to_vec()))
}
NaslSocket::Close => Err(FunctionErrorKind::WrongArgument(
NaslSocket::Closed => Err(FunctionErrorKind::WrongArgument(
"the given socket FD is already closed".to_string(),
)),
}
Expand Down Expand Up @@ -255,7 +255,7 @@ impl NaslSockets {
"This function is only available for TCP connections".to_string(),
None,
)),
NaslSocket::Close => Err(FunctionErrorKind::WrongArgument(
NaslSocket::Closed => Err(FunctionErrorKind::WrongArgument(
"the given socket FD is already closed".to_string(),
)),
}
Expand All @@ -267,13 +267,7 @@ impl NaslSockets {
/// - Secret/kdc_use_tcp
#[nasl_function]
fn open_sock_kdc(&self, context: &Context) -> Result<NaslValue, FunctionErrorKind> {
let hostname = match get_kb_item(context, "Secret/kdc_hostname")? {
Some(x) => Ok(x.to_string()),
None => Err(FunctionErrorKind::Diagnostic(
"KB key 'Secret/kdc_hostname' is not set".to_string(),
None,
)),
}?;
let hostname = get_kb_item_str(context, "Secret/kdc_hostname")?;

let ip = lookup_host(&hostname)
.map_err(|_| {
Expand Down Expand Up @@ -437,9 +431,10 @@ impl NaslSockets {
}
// Unsupported transport layer
None | Some(OpenvasEncaps::Max) => {
return Err(FunctionErrorKind::WrongArgument(format!(
"unsupported transport layer: {transport}(unknown)"
)))
return Err(FunctionErrorKind::Diagnostic(
format!("unsupported transport layer: {transport} (unknown)"),
None,
))
}
// TLS/SSL
Some(tls_version) => match tls_version {
Expand Down Expand Up @@ -472,9 +467,10 @@ impl NaslSockets {
}
}
_ => {
return Err(FunctionErrorKind::WrongArgument(format!(
"unsupported transport layer: {transport}{tls_version}"
)))
return Err(FunctionErrorKind::Diagnostic(
format!("unsupported transport layer: {transport} {tls_version}"),
None,
))
}
},
}
Expand All @@ -491,37 +487,10 @@ impl NaslSockets {
}

fn get_tls_conf(context: &Context) -> Result<TlsConfig, FunctionErrorKind> {
let cert_path = match get_kb_item(context, "Secret/tls_cert")? {
Some(x) => Ok(x.to_string()),
None => Err(FunctionErrorKind::Diagnostic(
"KB key 'Secret/tls_cert' is not set".to_string(),
None,
)),
}?;

let key_path = match get_kb_item(context, "Secret/tls_key")? {
Some(x) => Ok(x.to_string()),
None => Err(FunctionErrorKind::Diagnostic(
"KB key 'Secret/tls_key' is not set".to_string(),
None,
)),
}?;

let password = match get_kb_item(context, "Secret/tls_password")? {
Some(x) => Ok(x.to_string()),
None => Err(FunctionErrorKind::Diagnostic(
"KB key 'Secret/tls_password' is not set".to_string(),
None,
)),
}?;

let cafile_path = match get_kb_item(context, "Secret/tls_cafile")? {
Some(x) => Ok(x.to_string()),
None => Err(FunctionErrorKind::Diagnostic(
"KB key 'Secret/tls_cafile' is not set".to_string(),
None,
)),
}?;
let cert_path = get_kb_item_str(context, "Secret/tls_cert")?;
let key_path = get_kb_item_str(context, "Secret/tls_key")?;
let password = get_kb_item_str(context, "Secret/tls_password")?;
let cafile_path = get_kb_item_str(context, "Secret/cafile_path")?;

Ok(TlsConfig {
cert_path,
Expand Down Expand Up @@ -556,7 +525,7 @@ impl NaslSockets {
let port = match socket {
NaslSocket::Tcp(conn) => conn.local_addr()?.port(),
NaslSocket::Udp(conn) => conn.local_addr()?.port(),
NaslSocket::Close => {
NaslSocket::Closed => {
return Err(FunctionErrorKind::WrongArgument(
"the given socket FD is already closed".to_string(),
))
Expand Down Expand Up @@ -637,7 +606,7 @@ impl NaslSockets {
"This function is only available for TCP connections".to_string(),
None,
)),
NaslSocket::Close => Err(FunctionErrorKind::WrongArgument(
NaslSocket::Closed => Err(FunctionErrorKind::WrongArgument(
"the given socket FD is already closed".to_string(),
)),
}
Expand Down
33 changes: 16 additions & 17 deletions rust/src/nasl/builtin/network/udp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,30 +38,29 @@ impl Read for UdpConnection {
impl Write for UdpConnection {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let mtu = mtu(self.socket.peer_addr()?.ip());
if buf.len() < mtu {
let result = unsafe {
libc::send(
self.socket.as_raw_fd(),
buf.as_ptr() as *const libc::c_void,
buf.len(),
self.flags.unwrap_or_default(),
)
};
self.flags = None;
if result < 0 {
return Err(io::Error::last_os_error());
}
Ok(result as usize)
} else {
Err(io::Error::new(
if buf.len() > mtu {
return Err(io::Error::new(
io::ErrorKind::Other,
format!(
"UDP data of size {} exceeds the maximum length of {}",
buf.len(),
mtu
),
))
));
}
let result = unsafe {
libc::send(
self.socket.as_raw_fd(),
buf.as_ptr() as *const libc::c_void,
buf.len(),
self.flags.unwrap_or_default(),
)
};
self.flags = None;
if result < 0 {
return Err(io::Error::last_os_error());
}
Ok(result as usize)
}

fn flush(&mut self) -> io::Result<()> {
Expand Down

0 comments on commit 423af69

Please sign in to comment.