Merge pull request #39 from secure-software-engineering/meka

Add Multi-label ML approach using MEKA

Author: oshando

swan-pipeline/pom.xml

@@ -1,16 +1,16 @@
 package de.fraunhofer.iem.swan;
 
-import com.fasterxml.jackson.databind.ObjectMapper;
 import de.fraunhofer.iem.swan.cli.SwanOptions;
-import de.fraunhofer.iem.swan.features.FeaturesHandler;
+import de.fraunhofer.iem.swan.features.FeatureSetSelector;
+import de.fraunhofer.iem.swan.features.IFeatureSet;
 import de.fraunhofer.iem.swan.features.code.soot.SourceFileLoader;
 import de.fraunhofer.iem.swan.io.dataset.SrmList;
 import de.fraunhofer.iem.swan.io.dataset.SrmListUtils;
 import de.fraunhofer.iem.swan.model.ModelEvaluator;
 import de.fraunhofer.iem.swan.util.Util;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import java.io.File;
+
 import java.io.IOException;
 
 /**
@@ -48,19 +48,13 @@ public void run() throws IOException, InterruptedException {
         testDataset.load(dataset.getMethods());
 
         //Initialize and populate features
-        FeaturesHandler featuresHandler = new FeaturesHandler(dataset, testDataset, options);
-        featuresHandler.createFeatures();
+        FeatureSetSelector featureSetSelector = new FeatureSetSelector();
+        IFeatureSet featureSet = featureSetSelector.select(dataset, testDataset, options);
 
         //Train and evaluate model for SRM and CWE categories
-        ModelEvaluator modelEvaluator = new ModelEvaluator(featuresHandler, options);
+        ModelEvaluator modelEvaluator = new ModelEvaluator(featureSet, options, testDataset.getMethods());
         modelEvaluator.trainModel();
 
-        //TODO export final list to JSON file
-        String outputFile = options.getOutputDir() + File.separator + "swan-srm-cwe-list.json";
-        ObjectMapper objectMapper = new ObjectMapper();
-        objectMapper.writeValue(new File(outputFile), dataset);
-        logger.info("SRM/CWE list exported to {}", outputFile);
-
         long analysisTime = System.currentTimeMillis() - startAnalysisTime;
         logger.info("Total runtime {} minutes", analysisTime / 60000);
     }

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/SwanPipeline.java

@@ -25,8 +25,8 @@ public class CliRunner implements Callable<Integer> {
     @CommandLine.Option(names = {"-f", "--feature"}, description = {"Select one or more feature sets: all, code, doc-auto or doc-manual"})
     private List<String> featureSet =  Collections.singletonList("code");
 
-    @CommandLine.Option(names = {"-l", "--learning"}, description = {"Learning modes: manual, auto"})
-    private String learningMode = "manual";
+    @CommandLine.Option(names = {"-t", "--toolkit"}, description = {"ML toolkit: meka, weka, ml-plan"})
+    private String toolkit = "meka";
 
     @CommandLine.Option(names = {"-s", "--srm"}, description = {"SRM: all, source, sink, sanitizer, authentication, relevant"})
     private List<String> srmClasses = Collections.singletonList("all");
@@ -46,6 +46,13 @@ public class CliRunner implements Callable<Integer> {
     @CommandLine.Option(names = {"-sp", "--training-split"}, description = {"Percentage for training"})
     private double split = 0.7;
 
+    @CommandLine.Option(names = {"-p", "--phase"}, description = {"Phase: validate, predict"})
+    private String phase = "predict";
+
+    @CommandLine.Option(names = {"-pt", "--prediction-threshold"}, description = {"Threshold for predicting categories"})
+    private double predictionThreshold = 0.5;
+
+
     @Override
     public Integer call() throws Exception {
 
@@ -54,13 +61,15 @@ public Integer call() throws Exception {
                 datasetJson,
                 outputDir,
                 featureSet,
-                learningMode,
+                toolkit,
                 srmClasses,
                 cweClasses,
                 exportArffData,
                 isDocumented,
                 iterations,
-                split);
+                split,
+                phase);
+        options.setPredictionThreshold(predictionThreshold);
 
         return new SwanCli().run(options);
     }

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/cli/CliRunner.java

@@ -31,7 +31,7 @@ public Integer run(SwanOptions options) throws Exception {
         }
 
         if(options.getSrmClasses().contains("all")){
-            options.setSrmClasses(Arrays.asList("source", "sink", "sanitizer", "authentication", "relevant"));
+            options.setSrmClasses(Arrays.asList("source", "sink", "sanitizer", "authentication"));
         }
 
         if(options.getCweClasses().contains("all")){

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/cli/SwanCli.java

@@ -15,23 +15,44 @@ public class SwanOptions {
     private String datasetJson;
     private String outputDir;
     private List<String> featureSet;
-    private String learningMode;
+    private String toolkit;
     private List<String> srmClasses;
     private List<String> cweClasses;
     private boolean exportArffData;
     private boolean isDocumented;
     private int iterations;
     private double trainTestSplit;
+    private String phase;
+    private double predictionThreshold;
 
     public SwanOptions(String testDataDir, String trainDataDir, String datasetJson, String outputDir,
-                       List<String> featureSet, String learningMode, List<String> srmClasses, List<String> cweClasses,
+                       List<String> featureSet, String toolkit, List<String> srmClasses,
+                       List<String> cweClasses, boolean exportArffData, boolean isDocumented,
+                       int iterations, double trainTestSplit, String phase) {
+        this.testDataDir = testDataDir;
+        this.trainDataDir = trainDataDir;
+        this.datasetJson = datasetJson;
+        this.outputDir = outputDir;
+        this.featureSet = featureSet;
+        this.toolkit = toolkit;
+        this.srmClasses = srmClasses;
+        this.cweClasses = cweClasses;
+        this.exportArffData = exportArffData;
+        this.isDocumented = isDocumented;
+        this.iterations = iterations;
+        this.trainTestSplit = trainTestSplit;
+        this.phase = phase;
+    }
+
+    public SwanOptions(String testDataDir, String trainDataDir, String datasetJson, String outputDir,
+                       List<String> featureSet, String toolkit, List<String> srmClasses, List<String> cweClasses,
                        boolean exportArffData, boolean isDocumented, int iterations, double trainTestSplit) {
         this.testDataDir = testDataDir;
         this.trainDataDir = trainDataDir;
         this.datasetJson = datasetJson;
         this.outputDir = outputDir;
         this.featureSet = featureSet;
-        this.learningMode = learningMode;
+        this.toolkit = toolkit;
         this.srmClasses = srmClasses;
         this.cweClasses = cweClasses;
         this.exportArffData = exportArffData;
@@ -80,12 +101,12 @@ public void setFeatureSet(List<String> featureSet) {
         this.featureSet = featureSet;
     }
 
-    public String getLearningMode() {
-        return learningMode;
+    public String getToolkit() {
+        return toolkit;
     }
 
-    public void setLearningMode(String learningMode) {
-        this.learningMode = learningMode;
+    public void setToolkit(String toolkit) {
+        this.toolkit = toolkit;
     }
 
     public List<String> getSrmClasses() {
@@ -140,6 +161,22 @@ public void setTrainTestSplit(double trainTestSplit) {
         this.trainTestSplit = trainTestSplit;
     }
 
+    public String getPhase() {
+        return phase;
+    }
+
+    public void setPhase(String phase) {
+        this.phase = phase;
+    }
+
+    public double getPredictionThreshold() {
+        return predictionThreshold;
+    }
+
+    public void setPredictionThreshold(double predictionThreshold) {
+        this.predictionThreshold = predictionThreshold;
+    }
+
     @Override
     public String toString() {
         return "SwanOptions{" +
@@ -148,7 +185,7 @@ public String toString() {
                 ", datasetJson='" + datasetJson + '\'' +
                 ", outputDir='" + outputDir + '\'' +
                 ", featureSet='" + featureSet + '\'' +
-                ", learningMode='" + learningMode + '\'' +
+                ", learningMode='" + toolkit + '\'' +
                 ", srmClasses=" + srmClasses +
                 ", cweClasses=" + cweClasses +
                 ", exportArffData=" + exportArffData +

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/cli/SwanOptions.java

@@ -1,6 +1,7 @@
 
 package de.fraunhofer.iem.swan.data;
 
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 /**
@@ -42,9 +43,11 @@ public void setClassComment(String classComment) {
         this.classComment = classComment;
     }
 
+    @JsonIgnore
     public String getMergedComments() {
         return methodComment + " " + classComment;
     }
+
     @Override
     public String toString() {
         return "Javadoc{" +