Merge pull request #1239 from Privado-Inc/dev

Release PR
Privado-Inc · Aug 6, 2024 · 56febc6 · 56febc6
2 parents af3a8e5 + 7af75a0
commit 56febc6
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 13 deletions.
diff --git a/src/main/scala/ai/privado/languageEngine/java/language/package.scala b/src/main/scala/ai/privado/languageEngine/java/language/package.scala
@@ -25,26 +25,30 @@ package ai.privado.languageEngine.java
 
 import io.shiftleft.codepropertygraph.generated.{Cpg, EdgeTypes, NodeTypes}
 import io.shiftleft.codepropertygraph.generated.nodes.{AstNode, CfgNode, File, JavaProperty, Literal, MethodParameterIn}
-import overflowdb.traversal._
+import overflowdb.traversal.*
 
 import scala.jdk.CollectionConverters.IteratorHasAsScala
+import scala.util.Try
 
 package object language {
 
   implicit class NodeStarters(cpg: Cpg) {
     def property: Traversal[JavaProperty] =
-      cpg.graph.nodes(NodeTypes.JAVA_PROPERTY).asScala.cast[JavaProperty]
+      Try(cpg.graph.nodes(NodeTypes.JAVA_PROPERTY).asScala.cast[JavaProperty]).toOption
+        .getOrElse(Iterator.empty[JavaProperty])
   }
 
   implicit class StepsForProperty(val trav: Traversal[JavaProperty]) extends AnyVal {
 
-    def usedAt: Traversal[CfgNode] = trav.out(EdgeTypes.IS_USED_AT).cast[CfgNode]
-    def file: Traversal[File]      = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
+    def usedAt: Traversal[CfgNode] =
+      Try(trav.out(EdgeTypes.IS_USED_AT).cast[CfgNode]).toOption.getOrElse(Iterator.empty[CfgNode])
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
 
   }
 
   implicit class NodeTravToProperty(val trav: Traversal[AstNode]) {
-    def originalProperty: Traversal[JavaProperty] = trav.out(EdgeTypes.ORIGINAL_PROPERTY).cast[JavaProperty]
+    def originalProperty: Traversal[JavaProperty] =
+      Try(trav.out(EdgeTypes.ORIGINAL_PROPERTY).cast[JavaProperty]).toOption.getOrElse(Iterator.empty[JavaProperty])
   }
 
   implicit class NodeToProperty(val node: AstNode) {

diff --git a/src/main/scala/ai/privado/semantic/Language.scala b/src/main/scala/ai/privado/semantic/Language.scala
@@ -63,21 +63,23 @@ object Language {
   }
 
   implicit class StepsForPropertyForDbNode(val trav: Traversal[DbNode]) extends AnyVal {
-    def file: Traversal[File] = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
 
   }
   implicit class StepsForPropertyForSqlQueryNode(val trav: Traversal[SqlQueryNode]) extends AnyVal {
-    def file: Traversal[File]             = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
-    def sqlTable: Traversal[SqlTableNode] = trav.out(EdgeTypes.AST).cast[SqlTableNode]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
+    def sqlTable: Traversal[SqlTableNode] =
+      Try(trav.out(EdgeTypes.AST).cast[SqlTableNode]).toOption.getOrElse(Iterator.empty[SqlTableNode])
   }
 
   implicit class StepsForPropertyForSqlTableNode(val trav: Traversal[SqlTableNode]) extends AnyVal {
-    def file: Traversal[File]               = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
-    def sqlColumn: Traversal[SqlColumnNode] = trav.out(EdgeTypes.AST).cast[SqlColumnNode]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
+    def sqlColumn: Traversal[SqlColumnNode] =
+      Try(trav.out(EdgeTypes.AST).cast[SqlColumnNode]).toOption.getOrElse(Iterator.empty[SqlColumnNode])
   }
 
   implicit class StepsForPropertyForSqlColumnNode(val trav: Traversal[SqlColumnNode]) extends AnyVal {
-    def file: Traversal[File] = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
 
   }
 
@@ -97,12 +99,14 @@ object Language {
 
   implicit class NodeStarterForAndroidXmlLayoutNode(cpg: Cpg) {
     def androidXmlLayoutNode: Traversal[AndroidXmlLayoutNode] =
-      cpg.graph.nodes(NodeTypes.ANDROID_XML_LAYOUT_NODE).asScala.cast[AndroidXmlLayoutNode]
+      Try(cpg.graph.nodes(NodeTypes.ANDROID_XML_LAYOUT_NODE).asScala.cast[AndroidXmlLayoutNode]).toOption
+        .getOrElse(Iterator.empty[AndroidXmlLayoutNode])
   }
 
   implicit class NodeStarterForAndroidXmlPermissionNode(cpg: Cpg) {
     def androidXmlPermissionNode: Traversal[AndroidXmlPermissionNode] =
-      cpg.graph.nodes(NodeTypes.ANDROID_XML_PERMISSION_NODE).asScala.cast[AndroidXmlPermissionNode]
+      Try(cpg.graph.nodes(NodeTypes.ANDROID_XML_PERMISSION_NODE).asScala.cast[AndroidXmlPermissionNode]).toOption
+        .getOrElse(Iterator.empty[AndroidXmlPermissionNode])
   }
 
 }