Basic PoC for callgraph generation

examples/plugin/callgraph/main.go

@@ -0,0 +1,88 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+
+	"github.com/safedep/code/core"
+	"github.com/safedep/code/fs"
+	"github.com/safedep/code/lang"
+	"github.com/safedep/code/parser"
+	"github.com/safedep/code/plugin"
+	"github.com/safedep/code/plugin/callgraph"
+	"github.com/safedep/dry/log"
+)
+
+var (
+	dirToWalk string
+	language  string
+)
+
+func init() {
+	log.InitZapLogger("walker", "dev")
+
+	flag.StringVar(&dirToWalk, "dir", "", "Directory to walk")
+	flag.StringVar(&language, "lang", "python", "Language to use for parsing files")
+
+	flag.Parse()
+}
+
+func main() {
+	if dirToWalk == "" {
+		flag.Usage()
+		return
+	}
+
+	err := run()
+	if err != nil {
+		panic(err)
+	}
+}
+
+func run() error {
+	fileSystem, err := fs.NewLocalFileSystem(fs.LocalFileSystemConfig{
+		AppDirectories: []string{dirToWalk},
+	})
+
+	if err != nil {
+		return fmt.Errorf("failed to create local filesystem: %w", err)
+	}
+
+	language, err := lang.GetLanguage(language)
+	if err != nil {
+		return fmt.Errorf("failed to get language: %w", err)
+	}
+
+	walker, err := fs.NewSourceWalker(fs.SourceWalkerConfig{}, language)
+	if err != nil {
+		return fmt.Errorf("failed to create source walker: %w", err)
+	}
+
+	treeWalker, err := parser.NewWalkingParser(walker, language)
+	if err != nil {
+		return fmt.Errorf("failed to create tree walker: %w", err)
+	}
+
+	// consume callgraph
+	var callgraphCallback callgraph.CallgraphCallback = func(cg *callgraph.CallGraph) error {
+		cg.PrintCallGraph()
+
+		fmt.Println("DFS Traversal:")
+		for _, node := range cg.DFS() {
+			fmt.Println(node)
+		}
+		fmt.Println()
+		return nil
+	}
+
+	pluginExecutor, err := plugin.NewTreeWalkPluginExecutor(treeWalker, []core.Plugin{
+		callgraph.NewCallGraphPlugin(callgraphCallback),
+	})
+
+	if err != nil {
+		return fmt.Errorf("failed to create plugin executor: %w", err)
+	}
+
+	return pluginExecutor.Execute(context.Background(), fileSystem)
+}

examples/plugin/callgraph/test.py

@@ -0,0 +1,69 @@
+import base64
+from utils import printinit, printenc, printdec, printf2
+
+# Node must be generated, but shouldn't be part of DFS
+class EncodingUnused:
+    def __init__(self):
+        printinit("Initialized unused")
+        pass
+
+    def applyUnused(self, msg, func):
+        return func(msg)
+
+class Encoding:
+    def __init__(self):
+        printinit("Initialized")
+        pass
+
+    def apply(self, msg, func):
+        return func(msg)
+
+    # Unused
+    def apply2(self, msg, func):
+        return func(msg)
+
+encoder = Encoding()
+encoded = encoder.apply("Hello, World!".encode('utf-8'), base64.b64encode)
+printenc(encoded)
+decoded = encoder.apply(encoded, base64.b64decode)
+printdec(decoded)
+
+
+def f1(value):
+  f2(value)
+
+def f2(value):
+  printf2(value)
+  if value == 0:
+    return
+  f1(value-1)
+  pass
+
+def multiply(a, b):
+    return a * b
+
+f1(multiply(2, 3))
+
+def foo():
+  print("foo")
+  pass
+
+def bar():
+  print("bar")
+  pass
+
+def baz():
+  print("baz")
+  pass
+def useless():
+  print("useless")
+  baz()
+  pass
+
+xyz = foo
+
+print("GG")
+
+xyz = bar
+
+xyz()

examples/plugin/callgraph/testClass.py

@@ -0,0 +1,30 @@
+import base64
+from utils import printinit, printenc, printdec as pdec
+
+class Encoding:
+    def __init__(self):
+        printinit("Initialized")
+        pass
+
+    def apply(self, msg, func):
+        return func(msg)
+
+    # Unused
+    def apply2(self, msg, func):
+        return func(msg)
+
+def getenc():
+    return "encoded"
+
+encoder = Encoding()
+encoded = encoder.apply("Hello, World!".encode('utf-8'), base64.b64encode)
+printenc(encoded)
+decoded = encoder.apply(getenc(), base64.b64decode)
+pdec(decoded)
+
+class NoConstructorClass:
+    def show():
+        print("NoConstructorClass")
+        pass
+ncc = NoConstructorClass()
+ncc.show()

examples/plugin/callgraph/testFnAssignment.py

@@ -0,0 +1,42 @@
+import pprint
+from xyzprintmodule import xyzprint, xyzprint2, xyzprint3 as pxyz3
+
+customprintxyz = pxyz3
+customprintxyz = xyzprint2
+customprintxyz("GG")
+
+
+def foo():
+  pprint.pprint("foo")
+  pass
+
+def bar():
+  xyzprint("bar")
+  pass
+
+# unused
+def baz():
+  xyzprint2("baz")
+  pass
+
+xyz = foo
+print("GG")
+
+xyz = bar
+
+xyz() # current analysis will simulate both foo() & bar() calls
+
+def nestParent():
+  def nestChild():
+    xyzprint("nestChild")
+    def nestGrandChild():
+      xyzprint2("nestGrandChild")
+      pass
+    nestGrandChild()
+  nestChild()
+nestParent()
+
+def useless():
+  print("useless")
+  baz()
+  pass

examples/plugin/callgraph/testNestedFn.py

@@ -0,0 +1,38 @@
+import pprint
+from xyzprintmodule import xyzprint, xyzprint2
+from os import listdir as listdirfn, chmod
+
+def outerfn1():
+  chmod("outerfn1")
+  pass
+def outerfn2():
+  listdirfn("outerfn2")
+  pass
+
+def nestParent():
+  def parentScopedFn():
+    xyzprint("parentScopedFn")
+
+  def nestChild():
+    xyzprint("nestChild")
+    outerfn1()
+
+    def childScopedFn():
+      xyzprint("childScopedFn")
+
+    def nestGrandChildUseless():
+      xyzprint2("nestGrandChildUseless")
+
+    def nestGrandChild():
+      pprint.pp("nestGrandChild")
+      parentScopedFn()
+      outerfn2()
+      childScopedFn()
+
+    nestGrandChild()
+
+  outerfn1()
+  nestChild()
+
+nestParent()
+

examples/plugin/callgraph/testObjAssignment.py

@@ -0,0 +1,31 @@
+import pprint
+from xyz import printxyz as pxyz, printxyz2, printxyz3
+from os import listdir as listdirfn, chmod
+
+class ClassA:
+  def __init__(self):
+    pxyz("init")
+
+  def method1(self):
+    printxyz2("GG")
+
+
+class ClassB:
+  def __init__(self):
+    pxyz("init")
+
+  def method1(self):
+    printxyz2("GG")
+
+  def methodUnique(self):
+    printxyz3("GG")
+    pprint.pp("GG")
+
+
+def main():
+  x = ClassA()
+  x = ClassB()
+  y = x
+  y.method1()
+  y.methodUnique()
+main()

examples/plugin/callgraph/testScopes.py

@@ -0,0 +1,24 @@
+from pprint import pprint
+from xyzprintmodule import xyzprint as xprint, xyzprint, xyzprint2, xyzprint3
+
+def fn1():
+  xprint("very outer fn1")
+fn1()
+
+def fn2():
+  def fn1():
+    xyzprint("fn1 inside fn2")
+  fn1()
+
+  def fn3():
+    def fn4():
+      def fn1():
+        xyzprint3("fn1 inside fn4 inside fn3")
+      xyzprint2("fn4 inside fn3")
+      fn1() # must call fn1 inside fn4
+    fn1() # must call fn1 inside fn2
+    fn4()
+  fn3()
+
+fn2()
+

plugin/callgraph/assignment.go

@@ -0,0 +1,19 @@
+package callgraph
+
+type AssignmentGraph struct {
+	Assignments map[string][]string // Map of identifier to possible namespaces or other identifiers
+}
+
+func NewAssignmentGraph() *AssignmentGraph {
+	return &AssignmentGraph{Assignments: make(map[string][]string)}
+}
+
+// Add an assignment
+func (ag *AssignmentGraph) AddAssignment(identifier string, target string) {
+	ag.Assignments[identifier] = append(ag.Assignments[identifier], target)
+}
+
+// Resolve an identifier to its targets
+func (ag *AssignmentGraph) Resolve(identifier string) []string {
+	return ag.Assignments[identifier]
+}