fix #19 and #20

src/parser/extractor.rs

@@ -10,9 +10,7 @@ use lazy_static::lazy_static;
 use pep_508::{self, Spec};
 use regex::Regex;
 
-
-
-use toml::{de::Error, Value};
+use toml::{de::Error, Table, Value};
 
 pub fn extract_imports_python(text: String, imp: &mut Vec<Dependency>) {
     lazy_static! {
@@ -200,24 +198,35 @@ pub fn extract_imports_pyproject(
     // println!("{:#?}",toml_value);
 
     // Helper function to extract dependency values (version strings) including nested tables
-    fn extract_dependencies(table: &toml::value::Table, poetry: Option<bool>) -> Result<Vec<String>, Error> {
+    fn extract_dependencies(
+        table: &toml::value::Table,
+        poetry: Option<bool>,
+    ) -> Result<Vec<String>, Error> {
         let mut deps = Vec::new();
 
         // for [project] in pyproject.toml, the insides require a different sort of parsing
-        // for poetry you need both keys and values (as dependency name and version), 
+        // for poetry you need both keys and values (as dependency name and version),
         // for [project] the values are just enough and the keys are in the vec below
-        let projectlevel: Vec<&str> = vec!["dependencies", "optional-dependencies.docs"];
-
+        let projectlevel: Vec<&str> = vec![
+            "dependencies",
+            "optional-dependencies.docs",
+            "optional-dependencies",
+        ];
+
         for (key, version) in table {
             if projectlevel.contains(&key.as_str()) {
                 match version {
                     Value::String(version_str) => {
                         deps.push(version_str.to_string());
                     }
                     Value::Table(nested_table) => {
-                        // Recursively extract dependencies from nested tables
-                        let nested_deps = extract_dependencies(nested_table,None)?;
-                        deps.extend(nested_deps);
+                        if "optional-dependencies" == key {
+                            parse_opt_deps_pyproject(nested_table.clone(), &mut deps);
+                        } else {
+                            // Recursively extract dependencies from nested tables
+                            let nested_deps = extract_dependencies(nested_table, None)?;
+                            deps.extend(nested_deps);
+                        }
                     }
                     Value::Array(array) => {
                         // Extract dependencies from an array (if any)
@@ -229,34 +238,32 @@ pub fn extract_imports_pyproject(
                     }
                     _ => eprintln!("ERR: Invalid dependency syntax found while TOML parsing"),
                 }
-            }
-            else if poetry.unwrap_or(false) {
-                    match version {
-                        Value::String(version_str) => {
-                            let verstr = version_str.to_string();
-                            if verstr.contains('^') {
-                                let s = format!("{} >= {}", key, verstr.strip_prefix('^').unwrap());
-                                deps.push(s);
-                            }
-                            else if verstr == "*" {
+            } else if poetry.unwrap_or(false) {
+                match version {
+                    Value::String(version_str) => {
+                        let verstr = version_str.to_string();
+                        if verstr.contains('^') {
+                            let s = format!("{} >= {}", key, verstr.strip_prefix('^').unwrap());
+                            deps.push(s);
+                        } else if verstr == "*" {
                             deps.push(key.to_string());
-                            }
-                        }
-                        Value::Table(nested_table) => {
-                            // Recursively extract dependencies from nested tables
-                            let nested_deps = extract_dependencies(nested_table,None)?;
-                            deps.extend(nested_deps);
                         }
-                        Value::Array(array) => {
-                            // Extract dependencies from an array (if any)
-                            for item in array {
-                                if let Value::String(item_str) = item {
-                                    deps.push(item_str.to_string());
-                                }
+                    }
+                    Value::Table(nested_table) => {
+                        // Recursively extract dependencies from nested tables
+                        let nested_deps = extract_dependencies(nested_table, None)?;
+                        deps.extend(nested_deps);
+                    }
+                    Value::Array(array) => {
+                        // Extract dependencies from an array (if any)
+                        for item in array {
+                            if let Value::String(item_str) = item {
+                                deps.push(item_str.to_string());
                             }
                         }
-                        _ => eprintln!("ERR: Invalid dependency syntax found while TOML parsing"),
                     }
+                    _ => eprintln!("ERR: Invalid dependency syntax found while TOML parsing"),
+                }
             }
         }
         Ok(deps)
@@ -270,15 +277,15 @@ pub fn extract_imports_pyproject(
 
     for key in keys_to_check {
         if key.contains("tool") {
-
             if let Some(dependencies_table) = toml_value.get("tool") {
                 if let Some(dependencies_table) = dependencies_table.get("poetry") {
                     let poetrylevel: Vec<&str> = vec!["dependencies", "dev-dependencies"];
                     for k in poetrylevel.into_iter() {
                         if let Some(dep) = dependencies_table.get(k) {
                             match dep {
                                 Value::Table(table) => {
-                                    all_dependencies.extend(extract_dependencies(table, Some(true))?);                            
+                                    all_dependencies
+                                        .extend(extract_dependencies(table, Some(true))?);
                                 }
                                 // its definitely gonna be a table anyway, so...
                                 Value::String(_) => todo!(),
@@ -293,24 +300,24 @@ pub fn extract_imports_pyproject(
                 }
             }
         }
-
         // if its not poetry, check for [project] dependencies
         else if !key.contains("poetry") {
-
             if let Some(dependencies_table) = toml_value.get(key) {
                 if let Some(dependencies) = dependencies_table.as_table() {
-                        all_dependencies.extend(extract_dependencies(dependencies, None)?);
+                    all_dependencies.extend(extract_dependencies(dependencies, None)?);
                 }
             }
-        }
-        else {
-            eprintln!("The pyproject.toml seen here is unlike of a python project. Please check and make
-            sure you are in the right directory, or check the toml file."); exit(1)
+        } else {
+            eprintln!(
+                "The pyproject.toml seen here is unlike of a python project. Please check and make
+            sure you are in the right directory, or check the toml file."
+            );
+            exit(1)
         }
     }
-    // the toml might contain repeated dependencies 
+    // the toml might contain repeated dependencies
     // for different tools, dev tests, etc.
-    all_dependencies.dedup(); 
+    all_dependencies.dedup();
 
     for d in all_dependencies {
         let d = d.as_str();
@@ -350,3 +357,31 @@ pub fn extract_imports_pyproject(
     }
     Ok(())
 }
+
+pub fn parse_opt_deps_pyproject(table: Table, deps: &mut Vec<String>) {
+    for v in table.values() {
+        match v {
+            Value::Array(a) => {
+                for d in a {
+                    match d {
+                        Value::String(dependency) => {
+                            deps.push(dependency.to_owned());
+                        }
+                        Value::Integer(_) => todo!(),
+                        Value::Float(_) => todo!(),
+                        Value::Boolean(_) => todo!(),
+                        Value::Datetime(datetime) => todo!(),
+                        Value::Array(vec) => todo!(),
+                        Value::Table(map) => todo!(),
+                    }
+                }
+            }
+            Value::String(_) => todo!(),
+            Value::Integer(_) => todo!(),
+            Value::Float(_) => todo!(),
+            Value::Boolean(_) => todo!(),
+            Value::Datetime(datetime) => todo!(),
+            Value::Table(map) => todo!(),
+        }
+    }
+}

src/parser/mod.rs

@@ -14,6 +14,10 @@ pub async fn scan_dir(dir: &Path) {
     if let Ok(entries) = fs::read_dir(dir) {
         for entry in entries.flatten() {
             let filename = entry.file_name();
+            let filext = if let Some(ext) = Path::new(&filename).extension() {
+                ext.to_os_string()
+            } else {"none".into()};
+
 
             // setup.py check comes first otherwise it might cause issues with .py checker
             if *"setup.py" == filename.clone() {
@@ -25,8 +29,8 @@ pub async fn scan_dir(dir: &Path) {
                 result.setuppy();
             }
             // check if .py
-            // about the slice: [(file length) - 3..] for the extention
-            else if ".py" == &filename.to_str().unwrap()[{ filename.to_str().unwrap().len() - 3 }..] {
+            // checking file extension straight up from filename caused some bugs.
+            else if ".py" == filext {
                 result.add(FoundFile {
                     name: filename,
                     filetype: FileTypes::Python,

src/scanner/api.rs

@@ -135,7 +135,6 @@ impl Osv {
                     }
                 }
             }
-
             if let Ok(p) = parsed {
                 for vres in p.results {
                     if let Some(vulns) = vres.vulns {

src/scanner/mod.rs

@@ -18,7 +18,7 @@ pub async fn start(imports: Vec<Dependency>) -> Result<(), std::io::Error> {
     let collected = osv.query_batched(imports).await;
     // query_batched passes stuff onto display module after
 
-    // if everything went fine:
+    // if we collected vulns:
     if !collected.is_empty() {
         exit(1)
     }

src/scanner/models.rs

@@ -407,13 +407,21 @@ impl QueryBatched {
 
 impl Vulnerability {
     pub fn to_scanned_dependency(&self, imports_info: &HashMap<String, String>) -> ScannedDependency {
+        // println!("{:#?}", imports_info);
         let name_from_v = if let Some(n) = self.vulns.first() {
             if !n.affected.is_empty() {n.affected.first().unwrap().package.name.clone()}
             else {"Name in Context".to_string()}
         }
         else {"Name In Context".to_string()};
 
-        let version_from_map = imports_info.get(&name_from_v).unwrap(); // unwrapping safe as the hashmap is literally from the source of where the vuln was created...hopefully.
+        let version_from_map = if let Some(v) = imports_info.get(&name_from_v) {
+            v
+        } else {
+            &"parent package related to one of your dependencies".to_string()
+            // this happens rarely but every once in a while a vulnerability
+            // gets assigned to a different (usually) parent package so trying to
+            // get the version from our map fails unfortunately.
+        };
 
         ScannedDependency { name: name_from_v, version: version_from_map.to_owned(), vuln: self.clone() }