Add: nasl builtin functions for kb manipulation (#1703)

* use nasl function builtin macro for kb nasl functions

* Add: nasl builtin functions for kb manipulation
- replace_kb_item()
- get_kb_list()

rust/nasl-builtin-knowledge-base/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 nasl-builtin-utils = {path = "../nasl-builtin-utils"}
 nasl-syntax = {path = "../nasl-syntax"}
+nasl-function-proc-macro = {path = "../nasl-function-proc-macro"}
 storage = {path = "../storage"}
 
 [dev-dependencies]

rust/nasl-builtin-knowledge-base/README.md

@@ -2,9 +2,8 @@
 
 - set_kb_item
 - get_kp_item
-
-## Missing
-- get_host_kb_index
 - get_kb_list
-- index
 - replace_kb_item
+
+## Missing
+- get_host_kb_index: Do not apply. Redis specific and currently not used in any script

rust/nasl-builtin-knowledge-base/src/lib.rs

@@ -4,26 +4,32 @@
 
 use std::time::{SystemTime, UNIX_EPOCH};
 
-use nasl_builtin_utils::{error::FunctionErrorKind, get_named_parameter, NaslFunction};
+use nasl_builtin_utils::{error::FunctionErrorKind, NaslFunction};
 use storage::{Field, Kb, Retrieve};
 
 use nasl_builtin_utils::{Context, Register};
+use nasl_function_proc_macro::nasl_function;
 use nasl_syntax::NaslValue;
 
-/// NASL function to set a knowledge base
-fn set_kb_item(register: &Register, c: &Context) -> Result<NaslValue, FunctionErrorKind> {
-    let name = get_named_parameter(register, "name", true)?;
-    let value = get_named_parameter(register, "value", true)?;
-    let expires = match get_named_parameter(register, "expires", false) {
-        Ok(NaslValue::Number(x)) => Some(*x),
-        Ok(NaslValue::Exit(0)) => None,
-        Ok(x) => {
+/// NASL function to set a value under name in a knowledge base
+/// Only pushes unique values for the given name.
+#[nasl_function(named(name, value, expires))]
+fn set_kb_item(
+    name: NaslValue,
+    value: NaslValue,
+    expires: Option<NaslValue>,
+    c: &Context,
+) -> Result<NaslValue, FunctionErrorKind> {
+    let expires = match expires {
+        Some(NaslValue::Number(x)) => Some(x),
+        Some(NaslValue::Exit(0)) => None,
+        None => None,
+        Some(x) => {
             return Err(FunctionErrorKind::Diagnostic(
                 format!("expected expires to be a number but is {x}."),
                 None,
             ))
         }
-        Err(e) => return Err(e),
     }
     .map(|seconds| {
         let start = SystemTime::now();
@@ -46,33 +52,66 @@ fn set_kb_item(register: &Register, c: &Context) -> Result<NaslValue, FunctionEr
 }
 
 /// NASL function to get a knowledge base
-fn get_kb_item(register: &Register, c: &Context) -> Result<NaslValue, FunctionErrorKind> {
-    match register.positional() {
-        [x] => c
-            .retriever()
-            .retrieve(c.key(), Retrieve::KB(x.to_string()))
-            .map(|r| {
-                r.into_iter()
-                    .filter_map(|x| match x {
-                        Field::NVT(_) | Field::NotusAdvisory(_) | Field::Result(_) => None,
-                        Field::KB(kb) => Some(kb.value.into()),
-                    })
-                    .collect::<Vec<_>>()
-            })
-            .map(NaslValue::Fork)
-            .map_err(|e| e.into()),
-        x => Err(FunctionErrorKind::Diagnostic(
-            format!("expected one positional argument but got: {}", x.len()),
-            None,
-        )),
-    }
+#[nasl_function]
+fn get_kb_item(key: &str, c: &Context) -> Result<NaslValue, FunctionErrorKind> {
+    c.retriever()
+        .retrieve(c.key(), Retrieve::KB(key.to_string()))
+        .map(|r| {
+            r.into_iter()
+                .filter_map(|x| match x {
+                    Field::NVT(_) | Field::NotusAdvisory(_) | Field::Result(_) => None,
+                    Field::KB(kb) => Some(kb.value.into()),
+                })
+                .collect::<Vec<_>>()
+        })
+        .map(NaslValue::Fork)
+        .map_err(|e| e.into())
+}
+
+/// NASL function to replace a kb list
+#[nasl_function(named(name, value, expires))]
+fn replace_kb_item(
+    name: NaslValue,
+    value: NaslValue,
+    c: &Context,
+) -> Result<NaslValue, FunctionErrorKind> {
+    c.dispatcher()
+        .dispatch_replace(
+            c.key(),
+            Field::KB(Kb {
+                key: name.to_string(),
+                value: value.clone().as_primitive(),
+                expire: None,
+            }),
+        )
+        .map(|_| NaslValue::Null)
+        .map_err(|e| e.into())
+}
+
+/// NASL function to retrieve an item in a KB.
+#[nasl_function(named(name, value, expires))]
+fn get_kb_list(key: NaslValue, c: &Context) -> Result<NaslValue, FunctionErrorKind> {
+    c.retriever()
+        .retrieve(c.key(), Retrieve::KB(key.to_string()))
+        .map(|r| {
+            r.into_iter()
+                .filter_map(|x| match x {
+                    Field::NVT(_) | Field::NotusAdvisory(_) | Field::Result(_) => None,
+                    Field::KB(kb) => Some(kb.value.into()),
+                })
+                .collect::<Vec<_>>()
+        })
+        .map(NaslValue::Array)
+        .map_err(|e| e.into())
 }
 
 /// Returns found function for key or None when not found
 pub fn lookup(key: &str) -> Option<NaslFunction> {
     match key {
         "set_kb_item" => Some(set_kb_item),
         "get_kb_item" => Some(get_kb_item),
+        "get_kb_list" => Some(get_kb_list),
+        "replace_kb_item" => Some(replace_kb_item),
         _ => None,
     }
 }

rust/nasl-builtin-knowledge-base/tests/kb.rs

@@ -36,4 +36,41 @@ mod tests {
         assert_eq!(parser.next(), Some(Ok(NaslValue::Number(1))));
         assert!(matches!(parser.next(), Some(Err(_))));
     }
+    #[test]
+    fn get_kb_list() {
+        let code = r#"
+        set_kb_item(name: "test", value: 1);
+        set_kb_item(name: "test", value: 2);
+        get_kb_list("test");
+      
+        "#;
+        let register = Register::default();
+        let binding = ContextFactory::default();
+        let context = binding.build(Default::default(), Default::default());
+        let mut parser = CodeInterpreter::new(code, register, &context);
+        assert_eq!(parser.next(), Some(Ok(NaslValue::Null)));
+        assert_eq!(parser.next(), Some(Ok(NaslValue::Null)));
+        assert_eq!(
+            parser.next(),
+            Some(Ok(NaslValue::Array(vec![
+                NaslValue::Number(1),
+                NaslValue::Number(2)
+            ])))
+        );
+    }
+    #[test]
+    fn replace_kb_item() {
+        let code = r#"
+        set_kb_item(name: "test", value: 1);
+        replace_kb_item(name: "test", value: 2);
+        get_kb_item("test");
+        "#;
+        let register = Register::default();
+        let binding = ContextFactory::default();
+        let context = binding.build(Default::default(), Default::default());
+        let mut parser = CodeInterpreter::new(code, register, &context);
+        assert_eq!(parser.next(), Some(Ok(NaslValue::Null)));
+        assert_eq!(parser.next(), Some(Ok(NaslValue::Null)));
+        assert_eq!(parser.next(), Some(Ok(NaslValue::Number(2))));
+    }
 }

rust/openvasd/src/feed.rs

@@ -95,6 +95,14 @@ impl storage::Dispatcher for FeedIdentifier {
         Ok(())
     }
 
+    fn dispatch_replace(
+        &self,
+        _: &ContextKey,
+        _scope: storage::Field,
+    ) -> Result<(), storage::StorageError> {
+        Ok(())
+    }
+
     fn on_exit(&self) -> Result<(), storage::StorageError> {
         Ok(())
     }

rust/storage/src/item.rs

@@ -922,6 +922,15 @@ where
         }
     }
 
+    fn dispatch_replace(&self, key: &ContextKey, scope: Field) -> Result<(), StorageError> {
+        match scope {
+            Field::NVT(nvt) => self.store_nvt_field(nvt),
+            Field::KB(kb) => self.dispatcher.dispatch_kb(key, kb),
+            Field::NotusAdvisory(adv) => self.dispatcher.dispatch_advisory(key.as_ref(), *adv),
+            Field::Result(result) => self.dispatch(key, Field::Result(result)),
+        }
+    }
+
     fn on_exit(&self) -> Result<(), StorageError> {
         let mut data = Arc::as_ref(&self.nvt)
             .lock()

rust/storage/src/lib.rs

@@ -223,6 +223,10 @@ pub trait Dispatcher: Sync + Send {
     /// A key is usually a OID that was given when starting a script but in description run it is the filename.
     fn dispatch(&self, key: &ContextKey, scope: Field) -> Result<(), StorageError>;
 
+    /// Replace all fields under a key with the new field
+    ///
+    fn dispatch_replace(&self, key: &ContextKey, scope: Field) -> Result<(), StorageError>;
+
     /// On exit is called when a script exit
     ///
     /// Some database require a cleanup therefore this method is called when a script finishes.
@@ -256,6 +260,10 @@ where
         self.as_ref().dispatch(key, scope)
     }
 
+    fn dispatch_replace(&self, key: &ContextKey, scope: Field) -> Result<(), StorageError> {
+        self.as_ref().dispatch_replace(key, scope)
+    }
+
     fn on_exit(&self) -> Result<(), StorageError> {
         self.as_ref().on_exit()
     }
@@ -354,7 +362,25 @@ impl DefaultDispatcher {
         let mut data = self.kbs.as_ref().write()?;
         if let Some(scan_entry) = data.get_mut(scan_id) {
             if let Some(kb_entry) = scan_entry.get_mut(&kb.key) {
-                kb_entry.push(kb);
+                if kb_entry.iter().position(|x| x.value == kb.value).is_none() {
+                    kb_entry.push(kb);
+                };
+            } else {
+                scan_entry.insert(kb.key.clone(), vec![kb]);
+            }
+        } else {
+            let mut scan_entry = HashMap::new();
+            scan_entry.insert(kb.key.clone(), vec![kb]);
+            data.insert(scan_id.to_string(), scan_entry);
+        }
+        Ok(())
+    }
+
+    fn replace_kb(&self, scan_id: &str, kb: Kb) -> Result<(), StorageError> {
+        let mut data = self.kbs.as_ref().write()?;
+        if let Some(scan_entry) = data.get_mut(scan_id) {
+            if let Some(kb_entry) = scan_entry.get_mut(&kb.key) {
+                *kb_entry = vec![kb];
             } else {
                 scan_entry.insert(kb.key.clone(), vec![kb]);
             }
@@ -409,6 +435,20 @@ impl Dispatcher for DefaultDispatcher {
         Ok(())
     }
 
+    fn dispatch_replace(&self, key: &ContextKey, scope: Field) -> Result<(), StorageError> {
+        match scope {
+            Field::NVT(x) => self.cache_nvt_field(key.as_ref(), x)?,
+            Field::KB(x) => self.replace_kb(key.as_ref(), x)?,
+            Field::NotusAdvisory(x) => {
+                if let Some(x) = *x {
+                    self.cache_notus_advisory(x)?
+                }
+            }
+            Field::Result(x) => self.cache_result(key.as_ref(), *x)?,
+        }
+        Ok(())
+    }
+
     fn on_exit(&self) -> Result<(), StorageError> {
         if !self.dirty {
             self.cleanse()?;