add call graph

docker/cql.txt

@@ -41,5 +41,5 @@ CALL apoc.periodic.iterate("CALL apoc.load.csv('file:///Users/wh1t3P1g/Documents
 CALL apoc.periodic.iterate("CALL apoc.load.csv('file:///Users/wh1t3P1g/Documents/GitHub/tabby/docker/cache/interfaces.csv', {header:true}) YIELD map AS row RETURN row","MATCH (c1:Class {uuid: row.source}) MATCH (c2:Class {uuid: row.target}) MERGE (c1)-[e:INTERFACE {uuid: row.uuid}]->(c2)", {batchSize:1000, iterateList:true, parallel:true})
 CALL apoc.periodic.iterate("CALL apoc.load.csv('file:///Users/wh1t3P1g/Documents/GitHub/tabby/docker/cache/has.csv', {header:true}) YIELD map AS row RETURN row","MATCH (c:Class {uuid: row.classRef}) MATCH (m:Method {uuid: row.MethodRef}) MERGE (c)-[e:HAS {uuid: row.uuid}]->(m)", {batchSize:1000, iterateList:true, parallel:true})
 
-
+CALL apoc.periodic.iterate("CALL apoc.load.csv('file:///Users/wh1t3P1g/Documents/GitHub/tabby/docker/cache/calls.csv', {header:true}) YIELD map AS row RETURN row","MATCH (m1:Method {uuid: row.source}) MATCH (m2:Method {uuid: row.target}) MERGE (m1)-[e:CALL {uuid: row.uuid, lineNum: row.lineNum, realCallType: row.realCallType}]->(m2)", {batchSize:1000, iterateList:true, parallel:true})
 MATCH (c1:Class {uuid: row.source}) MATCH (c2:Class {uuid: row.target}) MERGE (c1)-[e:EXTENDS {uuid: row.uuid}]->(c2) RETURN *

src/main/java/tabby/config/GlobalConfiguration.java

@@ -26,7 +26,7 @@ public class GlobalConfiguration {
             new String[]{"uuid", "name", "signature", "isStatic", "hasParameters", "parameters", "returnType"},// method
             new String[]{"uuid", "source", "target"}, // extend/interfaces/
             new String[]{"uuid", "classRef", "MethodRef"}, // has
-            new String[]{"uuid", "source", "target", "lineNum"} // call
+            new String[]{"uuid", "source", "target", "lineNum", "realCallType"} // call
 
     ));
 

src/main/java/tabby/core/Analyser.java

@@ -5,12 +5,15 @@
 import org.springframework.stereotype.Component;
 import soot.*;
 import soot.options.Options;
+import tabby.config.GlobalConfiguration;
+import tabby.core.scanner.CallGraphScanner;
 import tabby.core.scanner.ClassInfoScanner;
 import tabby.dal.cache.CacheHelper;
 import tabby.util.FileUtils;
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.file.Files;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -28,6 +31,8 @@ public class Analyser {
     private CacheHelper cacheHelper;
     @Autowired
     private ClassInfoScanner classInfoScanner;
+    @Autowired
+    private CallGraphScanner callGraphScanner;
 
     /**
      * 运行当前soot分析
@@ -65,8 +70,11 @@ public void runSootAnalysis(String path, boolean isOnlyJDK){
                 Options.v().set_process_dir(getJdkDependencies());
                 cacheHelper.loadRuntimeClasses(getJdkDependencies(), true);
             }else{
-                Options.v().set_process_dir(FileUtils.getTargetDirectoryJarFiles(path));
-                cacheHelper.loadRuntimeClasses(FileUtils.getTargetDirectoryJarFiles(path), false);
+                List<String> targets = FileUtils.getTargetDirectoryJarFiles(path);
+                Options.v().set_soot_classpath(String.join(File.pathSeparator, getJdkDependencies()));
+                targets.addAll(getJdkDependencies());
+                Options.v().set_process_dir(targets);
+                cacheHelper.loadRuntimeClasses(targets, false);
             }
 
             Main.v().autoSetOptions();
@@ -78,9 +86,11 @@ public void runSootAnalysis(String path, boolean isOnlyJDK){
             // 类信息抽取
             classInfoScanner.run(cacheHelper.getRuntimeClasses());
             // 函数调用分析
+            PackManager.v().runPacks();
+            callGraphScanner.run(new ArrayList<>(cacheHelper.getSavedMethodRefs().values()));
             //            PhaseOptions.v().setPhaseOption("wjtp.classTransformer", "off");
-//            PackManager.v().runPacks();
-
+            classInfoScanner.save();
+            clean(); // clean caches
         }catch (CompilationDeathException e){
             if (e.getStatus() != CompilationDeathException.COMPILATION_SUCCEEDED) {
                 throw e;
@@ -101,4 +111,18 @@ public List<String> getJdkDependencies(){
         // TODO jdk其他的jar包是否也需要分析？
         return jdk;
     }
+
+    public void clean(){
+        try {
+            File cacheDir = new File(GlobalConfiguration.CACHE_PATH);
+            File[] files = cacheDir.listFiles();
+            if(files != null){
+                for(File file: files){
+                    Files.deleteIfExists(file.toPath());
+                }
+            }
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
 }

src/main/java/tabby/core/scanner/CallGraphScanner.java

@@ -0,0 +1,72 @@
+package tabby.core.scanner;
+
+import lombok.Data;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.stereotype.Component;
+import soot.SootMethod;
+import soot.Unit;
+import soot.jimple.InvokeExpr;
+import soot.jimple.JimpleBody;
+import soot.jimple.Stmt;
+import tabby.core.soot.switcher.InvokeStmtSwitcher;
+import tabby.dal.bean.ref.MethodReference;
+import tabby.dal.cache.CacheHelper;
+
+import java.util.List;
+
+/**
+ * 收集所有调用关系，这部分不做污点分析
+ * @author wh1t3P1g
+ * @since 2020/11/17
+ */
+@Data
+@Slf4j
+@Component
+public class CallGraphScanner implements Scanner<List<MethodReference>>{
+
+    @Autowired
+    private CacheHelper cacheHelper;
+    @Autowired
+    private InvokeStmtSwitcher invokeStmtSwitcher;
+
+    @Override
+    public void run(List<MethodReference> targets) {
+        collect(targets);
+        build();
+    }
+
+    @Override
+    public void collect(List<MethodReference> targets) {
+        log.info("start to build call graph!");
+        targets.forEach(this::collect);
+        log.info("build call graph DONE!");
+    }
+
+    public void collect(MethodReference methodRef){
+        try{
+            SootMethod method = methodRef.getCachedMethod();
+            JimpleBody body = (JimpleBody) method.getActiveBody();
+            invokeStmtSwitcher.setSource(methodRef);
+            for(Unit unit: body.getUnits()){
+                Stmt stmt = (Stmt) unit;
+                if(stmt.containsInvokeExpr()){
+                    InvokeExpr invokeExpr = stmt.getInvokeExpr();
+                    invokeExpr.apply(invokeStmtSwitcher);
+                }
+            }
+        }catch (RuntimeException e){
+//            e.printStackTrace();
+        }
+    }
+
+    @Override
+    public void build() {
+
+    }
+
+    @Override
+    public void save() {
+
+    }
+}

src/main/java/tabby/core/scanner/ClassInfoScanner.java

@@ -26,7 +26,7 @@
 @Data
 @Slf4j
 @Component
-public class ClassInfoScanner {
+public class ClassInfoScanner implements Scanner<List<String>> {
 
     @Autowired
     private CacheHelper cacheHelper;
@@ -35,12 +35,14 @@ public class ClassInfoScanner {
     @Autowired
     private MethodRefService methodRefService;
 
+    @Override
     public void run(List<String> classes){
         collect(classes);
         build();
-        save();
+//        save();
     }
 
+    @Override
     public void collect(List<String> classes){
         if(classes.isEmpty()) return;
 
@@ -49,6 +51,7 @@ public void collect(List<String> classes){
         log.info("collect "+classes.size()+" classes information. DONE!");
     }
 
+    @Override
     public void build(){
         if(cacheHelper.getSavedClassRefs().isEmpty()) return;
         Map<ClassRefHandle, ClassReference> clonedClassRefs = new HashMap<>(cacheHelper.getSavedClassRefs());
@@ -76,11 +79,12 @@ public void build(){
         });
     }
 
+    @Override
     public void save(){
         log.info("start to save cache to neo4j database!");
         // clear cache runtime classes
 //        cacheHelper.getRuntimeClasses().clear();
-        classRefService.clear();
+        classRefService.clear(); // TODO 初始化图数据库 正式版去掉
         // save cache to csv
         cacheHelper.saveToCSV();
         // load csv data to neo4j
@@ -92,6 +96,11 @@ public void save(){
         log.info("load csv data to neo4j finished!");
     }
 
+    /**
+     * 根据单个类进行类信息收集
+     * @param classname 待收集的类名
+     * @return 具体的类信息
+     */
     private ClassReference collect(String classname){
         ClassReference classRef = null;
         try{
@@ -105,7 +114,7 @@ private ClassReference collect(String classname){
             }
         }catch (Exception e){
             // class not found
-            log.debug(classname+" class not found!");
+//            log.debug(classname+" class not found!");
         }
 //        if(classRef == null){// 无法找到相应的类，只存储一个classname
         // // TODO 是否需要去存储没办法找到的类 存疑？
@@ -114,4 +123,6 @@ private ClassReference collect(String classname){
 //        }
         return classRef;
     }
+
+
 }

src/main/java/tabby/core/scanner/Scanner.java

@@ -0,0 +1,28 @@
+package tabby.core.scanner;
+
+/**
+ * @author wh1t3P1g
+ * @since 2020/11/17
+ */
+public interface Scanner<T> {
+
+
+    void run(T targets);
+
+    /**
+     * 收集类信息
+     * @param targets 待收集的类名/函数
+     */
+    void collect(T targets);
+
+    /**
+     * build relationships
+     */
+    void build();
+
+    /**
+     * save to cache file
+     * then cache file to neo4j
+     */
+    void save();
+}

src/main/java/tabby/core/soot/switcher/InvokeStmtSwitcher.java

@@ -1,11 +1,97 @@
 package tabby.core.soot.switcher;
 
-import soot.jimple.AbstractJimpleValueSwitch;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.stereotype.Component;
+import soot.SootMethodRef;
+import soot.Value;
+import soot.jimple.*;
+import soot.jimple.internal.JimpleLocal;
+import tabby.dal.bean.edge.Call;
+import tabby.dal.bean.ref.MethodReference;
+import tabby.dal.bean.ref.handle.ClassRefHandle;
+import tabby.dal.cache.CacheHelper;
+
+import java.util.List;
 
 /**
  * @author wh1t3P1g
  * @since 2020/10/10
  *
  */
+@Setter
+@Getter
+@Slf4j
+@Component
 public class InvokeStmtSwitcher extends AbstractJimpleValueSwitch {
+
+    private MethodReference source;
+
+    @Autowired
+    private CacheHelper cacheHelper;
+
+    @Override
+    public void caseStaticInvokeExpr(StaticInvokeExpr v) {
+        if(isNecessaryEdge("static", v)){
+            SootMethodRef sootMethodRef = v.getMethodRef();
+            ClassRefHandle classRefHandle = new ClassRefHandle(sootMethodRef.getDeclaringClass().getName());
+
+            buildCallRelationship(classRefHandle, sootMethodRef);
+        }
+    }
+
+    @Override
+    public void caseVirtualInvokeExpr(VirtualInvokeExpr v) {
+        SootMethodRef sootMethodRef = v.getMethodRef();
+        ClassRefHandle classRefHandle = new ClassRefHandle(v.getBase().getType().toString());
+        buildCallRelationship(classRefHandle, sootMethodRef);
+    }
+
+    @Override
+    public void caseSpecialInvokeExpr(SpecialInvokeExpr v) {// 初始化
+        SootMethodRef sootMethodRef = v.getMethodRef();
+        ClassRefHandle classRefHandle = new ClassRefHandle(v.getBase().getType().toString());
+        buildCallRelationship(classRefHandle, sootMethodRef);
+    }
+
+    @Override
+    public void caseInterfaceInvokeExpr(InterfaceInvokeExpr v) {
+        SootMethodRef sootMethodRef = v.getMethodRef();
+        ClassRefHandle classRefHandle = new ClassRefHandle(v.getBase().getType().toString());
+        buildCallRelationship(classRefHandle, sootMethodRef);
+    }
+
+    @Override
+    public void defaultCase(Object v) {
+        super.defaultCase(v);
+    }
+
+    public void buildCallRelationship(ClassRefHandle classRefHandle, SootMethodRef sootMethodRef){
+        MethodReference target = cacheHelper.loadMethodRef(classRefHandle, sootMethodRef.getName(), sootMethodRef.getSignature());
+        MethodReference source = cacheHelper.loadMethodRefByHandle(this.source.getHandle());
+        if(target != null && source != null){
+            Call call = Call.newInstance(source, target);
+            call.setRealCallType(classRefHandle.getName());
+            source.getCallEdge().add(call);
+        }
+    }
+
+    public <T> boolean isNecessaryEdge(String type, T v){
+        if ("static".equals(type)) { // 对于静态函数调用，只关注 函数参数可控的情况
+            StaticInvokeExpr invokeExpr = (StaticInvokeExpr) v;
+            if (invokeExpr.getArgCount() == 0) {
+                return false;
+            }
+            List<Value> values = invokeExpr.getArgs();
+            for (Value value : values) {
+                if (value instanceof JimpleLocal ||
+                        value instanceof StringConstant) { // Class.forName(xxx) 这种情况
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
 }

src/main/java/tabby/core/soot/toolkit/SimpleCallGraphExtractor.java

@@ -0,0 +1,8 @@
+package tabby.core.soot.toolkit;
+
+/**
+ * @author wh1t3P1g
+ * @since 2020/11/17
+ */
+public class SimpleCallGraphExtractor {
+}

src/main/java/tabby/core/soot/transformer/CallGraphTransformer.java

@@ -0,0 +1,17 @@
+package tabby.core.soot.transformer;
+
+import soot.Body;
+import soot.BodyTransformer;
+
+import java.util.Map;
+
+/**
+ * @author wh1t3P1g
+ * @since 2020/11/9
+ */
+public class CallGraphTransformer extends BodyTransformer {
+    @Override
+    protected void internalTransform(Body b, String phaseName, Map<String, String> options) {
+
+    }
+}