<commit_before>package magnum.antimalware

import scala.collection.mutable.ListBuffer
import scala.util.Random

import scala.sys.process._
import scala.language.postfixOps

import scala.io.Source

import java.io.{File,FileWriter}

import weka.core.Instances
import weka.core.converters.ArffLoader
import weka.classifiers.{Classifier,Evaluation}

// Classifiers.
import weka.classifiers.bayes.{BayesNet,NaiveBayes}
import weka.classifiers.functions.{Logistic,MultilayerPerceptron}
import weka.classifiers.trees.{J48,RandomForest}

class WekaWrapper(arffTag: String, resultsTag: String,
    percentBenignTrain: Int, percentMaliciousTrain: Int, randomSeed: Long) {
  def run() {
    Random.setSeed(randomSeed)

    val resultsDirStr = s"../results/$resultsTag"
    val resultsDir = new File(resultsDirStr)
    org.apache.commons.io.FileUtils.deleteDirectory(resultsDir) 
    resultsDir.mkdir;

    // Get lists of benign and malicious arffs.
    val (benignTrainingArffs, benignTestingArffs) =
      getArffs("benign", percentBenignTrain)
    val (maliciousTrainingArffs, maliciousTestingArffs) =
      getArffs("malicious", percentMaliciousTrain)
    val trainingArffs = benignTrainingArffs ++ maliciousTrainingArffs
    val testingArffs = benignTestingArffs ++ maliciousTestingArffs
    val trainingCombinedArff = combineAllArffs(trainingArffs, "training")
    val testingCombinedArff = combineAllArffs(testingArffs, "testing")

    val analysisFile = new FileWriter(s"$resultsDirStr/arff-analysis.txt")
    analysisFile.write("ARFF analysis.\n\n")
    analysisFile.write(s"arffTag: $arffTag\n")
    analysisFile.write(s"resultsTag: $resultsTag\n")
    analysisFile.write(s"percentBenignTrain: $percentBenignTrain\n")
    analysisFile.write(s"percentMaliciousTrain: $percentMaliciousTrain\n\n")
    dumpInfo(benignTrainingArffs ++ benignTestingArffs, analysisFile, "benign")
    dumpInfo(maliciousTrainingArffs ++ maliciousTestingArffs, analysisFile,
      "malicious")
    dumpInfo(benignTrainingArffs ++ benignTestingArffs ++
      maliciousTrainingArffs ++ maliciousTestingArffs, analysisFile, "all")
    analysisFile.write(s"Dump of `../arff/$arffTag/info.txt`.\n\n")
    for(line <- Source.fromFile(s"../arff/$arffTag/info.txt").getLines())
      analysisFile.write(line + "\n")
    analysisFile.close


    var arffLoader = new ArffLoader
    arffLoader.setFile(trainingCombinedArff)
    val train = arffLoader.getDataSet
    train.setClassIndex(train.numAttributes-1)

    /*
    // Let's remove training completely.
    // This pollutes the training set with testing feature vectors.
    // TODO: Hamilton, confirm we don't need this.
    if (mode == "Training") {
      val cvDir = new File("Results/CV")
      org.apache.commons.io.FileUtils.deleteDirectory(cvDir) 
      cvDir.mkdir

      // TODO Do these need any options?
      crossValidate(new BayesNet, "BayesNet", train)
      crossValidate(new NaiveBayes, "NaiveBayes", train)
      crossValidate(new Logistic, "Logistic", train)
      crossValidate(new MultilayerPerceptron, "MultilayerPerceptron",
                    train)
      crossValidate(new J48, "J48", train)
      crossValidate(new RandomForest, "RandomForest", train)
    } else {
      */
    val ttStr = s"$resultsDir/traintest"
    val ttDir = new File(ttStr)
    org.apache.commons.io.FileUtils.deleteDirectory(ttDir) 
    ttDir.mkdir;

    arffLoader.reset
    arffLoader.setFile(testingCombinedArff)
    val test = arffLoader.getDataSet
    test.setClassIndex(test.numAttributes-1)
    collectiveTest(new BayesNet, "BayesNet", train, test, ttStr)
    collectiveTest(new NaiveBayes, "NaiveBayes", train, test, ttStr)
    collectiveTest(new Logistic, "Logistic", train, test, ttStr)
    collectiveTest(new MultilayerPerceptron, "MultilayerPerceptron",
                   train, test, ttStr)
    collectiveTest(new J48, "J48", train, test, ttStr)
    collectiveTest(new RandomForest, "RandomForest", train, test, ttStr)

    individualTest(new BayesNet, "BayesNet", train, testingArffs, ttStr)
    individualTest(new NaiveBayes, "NaiveBayes", train, testingArffs, ttStr)
    individualTest(new Logistic, "Logistic", train, testingArffs, ttStr)
    individualTest(new MultilayerPerceptron, "MultilayerPerceptron",
                   train, testingArffs, ttStr)
    individualTest(new J48, "J48", train, testingArffs, ttStr)
    individualTest(new RandomForest, "RandomForest", train, testingArffs, ttStr)
  }

  def getArffs(mode: String, percentTrain: Int):
      (ListBuffer[File], ListBuffer[File]) = {
    val arffDir: File = new File(s"../arff/$arffTag/$mode")
    val arffList = arffDir.listFiles

    if (arffList == null) {
      print(s"Warning: Arff list for '$mode' is empty.")
      return (ListBuffer[File](), ListBuffer[File]())
    }

    val shuffledIndices = util.Random.shuffle(List.range(0, arffList.length))
    var trainingArffs: ListBuffer[File] = new ListBuffer[File]()
    var testingArffs: ListBuffer[File] = new ListBuffer[File]()

    var i = 0;
    while (i < (percentTrain/100.0) * arffList.length) {
      trainingArffs += arffList(shuffledIndices(i))
      i += 1
    }
    while (i < arffList.length) {
      testingArffs += arffList(shuffledIndices(i))
      i += 1
    }

    return (trainingArffs, testingArffs)
  }

  def dumpInfo(files: ListBuffer[File], fw: FileWriter, tag: String) {
    var arffLoader = new ArffLoader

    var fileCount = 0; var vectorCount = 0;
    for (arffFile <- files) {
      arffLoader.setFile(arffFile)
      fileCount += 1
      vectorCount += arffLoader.getDataSet.numInstances
      arffLoader.reset
    }

    fw.write(s"For $tag apps:\n")
    fw.write(s"  + Number of files: $fileCount\n")
    fw.write(s"  + Number of feature vectors: $vectorCount\n")
    fw.write(s"  + FV/File: ${vectorCount.toDouble/fileCount.toDouble}\n\n")
  }

  def combineAllArffs(files: ListBuffer[File], mode: String): File = {
    var arffLoader = new ArffLoader<commit_msg>Antimalware: Filter ARFF files so we don't pass non-ARFF files to Weka.<commit_after>package magnum.antimalware

import scala.collection.mutable.ListBuffer
import scala.util.Random

import scala.sys.process._
import scala.language.postfixOps

import scala.io.Source

import java.io.{File,FileWriter}

import weka.core.Instances
import weka.core.converters.ArffLoader
import weka.classifiers.{Classifier,Evaluation}

// Classifiers.
import weka.classifiers.bayes.{BayesNet,NaiveBayes}
import weka.classifiers.functions.{Logistic,MultilayerPerceptron}
import weka.classifiers.trees.{J48,RandomForest}

class WekaWrapper(arffTag: String, resultsTag: String,
    percentBenignTrain: Int, percentMaliciousTrain: Int, randomSeed: Long) {
  def run() {
    Random.setSeed(randomSeed)

    val resultsDirStr = s"../results/$resultsTag"
    val resultsDir = new File(resultsDirStr)
    org.apache.commons.io.FileUtils.deleteDirectory(resultsDir) 
    resultsDir.mkdir;

    // Get lists of benign and malicious arffs.
    val (benignTrainingArffs, benignTestingArffs) =
      getArffs("benign", percentBenignTrain)
    val (maliciousTrainingArffs, maliciousTestingArffs) =
      getArffs("malicious", percentMaliciousTrain)
    val trainingArffs = benignTrainingArffs ++ maliciousTrainingArffs
    val testingArffs = benignTestingArffs ++ maliciousTestingArffs
    val trainingCombinedArff = combineAllArffs(trainingArffs, "training")
    val testingCombinedArff = combineAllArffs(testingArffs, "testing")

    val analysisFile = new FileWriter(s"$resultsDirStr/arff-analysis.txt")
    analysisFile.write("ARFF analysis.\n\n")
    analysisFile.write(s"arffTag: $arffTag\n")
    analysisFile.write(s"resultsTag: $resultsTag\n")
    analysisFile.write(s"percentBenignTrain: $percentBenignTrain\n")
    analysisFile.write(s"percentMaliciousTrain: $percentMaliciousTrain\n\n")
    dumpInfo(benignTrainingArffs ++ benignTestingArffs, analysisFile, "benign")
    dumpInfo(maliciousTrainingArffs ++ maliciousTestingArffs, analysisFile,
      "malicious")
    dumpInfo(benignTrainingArffs ++ benignTestingArffs ++
      maliciousTrainingArffs ++ maliciousTestingArffs, analysisFile, "all")
    analysisFile.write(s"Dump of `../arff/$arffTag/info.txt`.\n\n")
    for(line <- Source.fromFile(s"../arff/$arffTag/info.txt").getLines())
      analysisFile.write(line + "\n")
    analysisFile.close


    var arffLoader = new ArffLoader
    arffLoader.setFile(trainingCombinedArff)
    val train = arffLoader.getDataSet
    train.setClassIndex(train.numAttributes-1)

    /*
    // Let's remove training completely.
    // This pollutes the training set with testing feature vectors.
    // TODO: Hamilton, confirm we don't need this.
    if (mode == "Training") {
      val cvDir = new File("Results/CV")
      org.apache.commons.io.FileUtils.deleteDirectory(cvDir) 
      cvDir.mkdir

      // TODO Do these need any options?
      crossValidate(new BayesNet, "BayesNet", train)
      crossValidate(new NaiveBayes, "NaiveBayes", train)
      crossValidate(new Logistic, "Logistic", train)
      crossValidate(new MultilayerPerceptron, "MultilayerPerceptron",
                    train)
      crossValidate(new J48, "J48", train)
      crossValidate(new RandomForest, "RandomForest", train)
    } else {
      */
    val ttStr = s"$resultsDir/traintest"
    val ttDir = new File(ttStr)
    org.apache.commons.io.FileUtils.deleteDirectory(ttDir) 
    ttDir.mkdir;

    arffLoader.reset
    arffLoader.setFile(testingCombinedArff)
    val test = arffLoader.getDataSet
    test.setClassIndex(test.numAttributes-1)
    collectiveTest(new BayesNet, "BayesNet", train, test, ttStr)
    collectiveTest(new NaiveBayes, "NaiveBayes", train, test, ttStr)
    collectiveTest(new Logistic, "Logistic", train, test, ttStr)
    collectiveTest(new MultilayerPerceptron, "MultilayerPerceptron",
                   train, test, ttStr)
    collectiveTest(new J48, "J48", train, test, ttStr)
    collectiveTest(new RandomForest, "RandomForest", train, test, ttStr)

    individualTest(new BayesNet, "BayesNet", train, testingArffs, ttStr)
    individualTest(new NaiveBayes, "NaiveBayes", train, testingArffs, ttStr)
    individualTest(new Logistic, "Logistic", train, testingArffs, ttStr)
    individualTest(new MultilayerPerceptron, "MultilayerPerceptron",
                   train, testingArffs, ttStr)
    individualTest(new J48, "J48", train, testingArffs, ttStr)
    individualTest(new RandomForest, "RandomForest", train, testingArffs, ttStr)
  }

  def getArffs(mode: String, percentTrain: Int):
      (ListBuffer[File], ListBuffer[File]) = {
    val arffDir: File = new File(s"../arff/$arffTag/$mode")
    val unfilteredArffList = arffDir.listFiles

    if (unfilteredArffList == null) {
      print(s"Warning: Arff list for '$mode' is empty.")
      return (ListBuffer[File](), ListBuffer[File]())
    }

    var arffList = new ListBuffer[File]
    var i = 0
    while (i < unfilteredArffList.length) {
      if (unfilteredArffList(i).getName().endsWith(".arff")) {
        arffList += unfilteredArffList(i);
      }
      i += 1
    }

    val shuffledIndices = util.Random.shuffle(List.range(0, arffList.length))
    var trainingArffs: ListBuffer[File] = new ListBuffer[File]()
    var testingArffs: ListBuffer[File] = new ListBuffer[File]()

    i = 0;
    while (i < (percentTrain/100.0) * arffList.length) {
      trainingArffs += arffList(shuffledIndices(i))
      i += 1
    }
    while (i < arffList.length) {
      testingArffs += arffList(shuffledIndices(i))
      i += 1
    }

    return (trainingArffs, testingArffs)
  }

  def dumpInfo(files: ListBuffer[File], fw: FileWriter, tag: String) {
    var arffLoader = new ArffLoader

    var fileCount = 0; var vectorCount = 0;
    for (arffFile <- files) {
      arffLoader.setFile(arffFile)
      fileCount += 1
      vectorCount += arffLoader.getDataSet.numInstances
      arffLoader.reset
    }

    fw.write(s"For $tag apps:\n")
    fw.write(s"  + Number of files: $fileCount\n")
    fw.write(s"  + Number of feature vectors: $vectorCount\n")
    fw.write(s"  + FV/File: ${vectorCount.toDouble/fileCount.toDouble}\n\n")
  }

  def combineAllArffs(files: ListBuffer[File], mode: String): File = {
    var arffLoader = new ArffLoader
{'fraction_non_alphanumeric': 0.11292962356792144, 'fraction_numerical': 0.004582651391162029, 'mean_word_length': 3.9258363563079404, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6892708', 'n_tokens_mistral': 4163, 'n_tokens_neox': 4011, 'n_words': 807}