io.shiftleft.SerializedCpg.scala Maven / Gradle / Ivy
package io.shiftleft
import java.io.{File, IOException}
import java.net.{URI, URISyntaxException}
import java.nio.file.{FileSystem, FileSystems, Files}
import java.util
import com.google.protobuf.GeneratedMessageV3
class SerializedCpg extends AutoCloseable {
/**
* We allow creating a dummy serialized CPG that does not do anything.
*/
private[this] var zipFileSystem: FileSystem = null
private[this] var counter = 0
/**
* Create Serialized CPG from existing file. If the file does not exist,
* an empty Serialized CPG is created.
**/
def this(filename: String) {
this()
initZipFilesystem(filename)
}
def isEmpty: Boolean = zipFileSystem == null
@throws[URISyntaxException]
@throws[IOException]
private[this] def initZipFilesystem(filename: String): Unit = {
val env = new util.HashMap[String, AnyRef]
// This ensures that the file is created if it does not exist
env.put("create", "true")
env.put("useTempFile", java.lang.Boolean.TRUE)
val fileUri = new File(filename).toURI
val outputUri = new URI("jar:" + fileUri.getScheme, null, fileUri.getPath, null)
zipFileSystem = FileSystems.newFileSystem(outputUri, env)
}
/**
* Add overlay graph named `name` to the zip file
**/
@throws[IOException]
def addOverlay(overlay: GeneratedMessageV3, name: String): Unit = {
if (!isEmpty) {
val pathInZip = zipFileSystem.getPath(s"${counter}_${name}")
counter += 1
val outputStream = Files.newOutputStream(pathInZip)
overlay.writeTo(outputStream)
outputStream.close()
}
}
@throws[IOException]
def addOverlay(overlays: Iterator[GeneratedMessageV3], name: String): Unit = {
overlays.zipWithIndex.foreach {
case (overlay, i) => addOverlay(overlay, name + "_" + i)
}
}
@throws[IOException]
override def close(): Unit = {
if (!isEmpty) {
zipFileSystem.close()
}
}
}