Merge pull request #18 from rohaquinlop/issue-16

feat(build): #16 enhance try cognitive complexity
rohaquinlop · Mar 4, 2024 · 4e64951 · 4e64951
2 parents b8bba38 + f738c72
commit 4e64951
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 631 deletions.
diff --git a/src/cognitive_complexity/mod.rs b/src/cognitive_complexity/mod.rs
@@ -153,12 +153,28 @@ fn statement_cognitive_complexity(statement: Stmt, nesting_level: u64) -> PyResu
     match statement {
         Stmt::FunctionDef(f) => {
             for node in f.body.iter() {
-                complexity += statement_cognitive_complexity(node.clone(), nesting_level)?;
+                match node {
+                    Stmt::FunctionDef(..) | Stmt::AsyncFunctionDef(..) => {
+                        complexity +=
+                            statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
+                    }
+                    _ => {
+                        complexity += statement_cognitive_complexity(node.clone(), nesting_level)?;
+                    }
+                }
             }
         }
         Stmt::AsyncFunctionDef(f) => {
             for node in f.body.iter() {
-                complexity += statement_cognitive_complexity(node.clone(), nesting_level)?;
+                match node {
+                    Stmt::FunctionDef(..) | Stmt::AsyncFunctionDef(..) => {
+                        complexity +=
+                            statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
+                    }
+                    _ => {
+                        complexity += statement_cognitive_complexity(node.clone(), nesting_level)?;
+                    }
+                }
             }
         }
         Stmt::ClassDef(c) => {
@@ -168,23 +184,26 @@ fn statement_cognitive_complexity(statement: Stmt, nesting_level: u64) -> PyResu
         }
         Stmt::Assign(a) => {
             complexity += count_bool_ops(*a.value);
+
+            if complexity > 0 {
+                complexity += nesting_level;
+            }
         }
-        // breaks in the linear flow
         Stmt::For(f) => {
-            complexity += 1;
+            complexity += 1 + nesting_level;
             for node in f.body.iter() {
                 complexity += statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
             }
         }
         Stmt::While(w) => {
-            complexity += 1;
+            complexity += 1 + nesting_level;
             complexity += count_bool_ops(*w.test);
             for node in w.body.iter() {
                 complexity += statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
             }
         }
         Stmt::If(i) => {
-            complexity += 1;
+            complexity += 1 + nesting_level;
             complexity += count_bool_ops(*i.test);
             for node in i.body.iter() {
                 complexity += statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
@@ -203,12 +222,12 @@ fn statement_cognitive_complexity(statement: Stmt, nesting_level: u64) -> PyResu
             }
         }
         Stmt::Try(t) => {
-            complexity += 1;
             for node in t.body.iter() {
                 complexity += statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
             }
 
             for handler in t.handlers.iter() {
+                complexity += 1;
                 match handler {
                     ast::ExceptHandler::ExceptHandler(e) => {
                         for node in e.body.iter() {
@@ -218,20 +237,32 @@ fn statement_cognitive_complexity(statement: Stmt, nesting_level: u64) -> PyResu
                     }
                 }
             }
+
+            for node in t.orelse.iter() {
+                complexity += statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
+            }
+
+            for node in t.finalbody.iter() {
+                complexity += statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
+            }
+
+            if complexity > 0 {
+                complexity += nesting_level;
+            }
         }
         Stmt::Match(m) => {
             for case in m.cases.iter() {
                 for node in case.body.iter() {
                     complexity += statement_cognitive_complexity(node.clone(), nesting_level + 1)?;
                 }
             }
+
+            if complexity > 0 {
+                complexity += nesting_level;
+            }
         }
         _ => {}
     };
 
-    if complexity > 0 {
-        complexity += nesting_level;
-    }
-
     Ok(complexity)
 }
diff --git a/tests/test.py b/tests/test.py
@@ -0,0 +1,30 @@
+# code taken from:
+# https://gitlab.com/fluidattacks/universe/-/blob/trunk/skims/skims/lib_sast/root/f267/kubernetes.py
+
+
+@atatus.capture_span()
+def k8s_check_host_pid(
+    shard: GraphShard, method_supplies: MethodSupplies
+) -> MethodExecutionResult:
+    method = MethodsEnum.K8S_CHECK_HOST_PID
+
+    def n_ids() -> Iterator[NId]:
+        graph = shard.syntax_graph
+        integrity = check_template_integrity_for_all_nodes(
+            graph, method_supplies.selected_nodes
+        )
+        if integrity:  # 2 (nested = 1)
+            for nid in method_supplies.selected_nodes:  # 3 (nested = 2)
+                for c_id in get_host_pid(graph, nid):  # 4 (nested = 3)
+                    yield c_id
+
+    return get_vulnerabilities_from_n_ids(
+        n_ids=n_ids(),
+        query_supplies=QuerySupplies(
+            desc_key="f267.k8s_check_host_pid",
+            desc_params={},
+            method=method,
+            method_calls=len(method_supplies.selected_nodes),
+            graph_shard=shard,
+        ),
+    )