<commit_before>#!/usr/bin/python3

# A Python script to generate test data given ELF or PE files including their
# relevant debug info.
#
# This Python script is meant to replace the horrible bash script that is
# currently responsible for doing this.

import subprocess, os, fnmatch, codecs, random, shutil
from subprocess import Popen, PIPE, STDOUT
from operator import itemgetter
from collections import defaultdict

work_directory="/var/tmp/train2"
shutil.rmtree(work_directory)
os.mkdir(work_directory)

def FindELFTrainingFiles():
  """ Returns the list of ELF files that should be used for training. These
  ELF files need to contain objdump-able debug information.

  """
  dirnames_and_masks = [ ('./unrar.5.5.3.builds', 'unrar.x??.O?') ]

  filenames = []
  for dirname, mask in dirnames_and_masks:
    filenames = filenames + [
      dirname + "/" + filename for filename in os.listdir(dirname)
      if fnmatch.fnmatch(filename, mask) ]
  return filenames

def FindPETrainingFiles():
  """ Returns the list of PE files that should be used for training. These
  PE files need to have associated text files (with suffix .debugdump) that
  contains the output of dia2dump in the same directory. """
  return [
   "./unrar.5.5.3.builds/VS2015/unrar32/Release/UnRAR.exe",
   "./unrar.5.5.3.builds/VS2015/unrar32/Debug/UnRAR.exe",
   "./unrar.5.5.3.builds/VS2015/unrar64/Release/UnRAR.exe",
   "./unrar.5.5.3.builds/VS2015/unrar64/Debug/UnRAR.exe" ]

def ObtainFunctionSymbols(training_file, file_format):
  if file_format == "ELF":
    return ObtainELFFunctionSymbols(training_file)
  elif file_format == "PE":
    return ObtainPEFunctionSymbols(training_file)

def SaneBase64(input_string):
  """ Because Python3 attempts to win 'most idiotic language ever', encoding a
  simple string as base64 without risking to have strange newlines added is
  difficult. This functions is an insane solution: Call command line
  base64encode instead of dealing with Python. """
  encoded_string = subprocess.run(["base64", "-w0"], stdout=PIPE,
    input=bytes(input_string, encoding="utf-8")).stdout.decode("utf-8")
  return encoded_string

def ObtainELFFunctionSymbols(training_file):
  """ Runs objdump to obtain the symbols in an ELF file and then returns a
  dictionary for this file. """
  result = {}
  symbols = [ line for line in subprocess.check_output(
    [ "objdump", "-t", training_file ] ).decode("utf-8").split("\n")
      if line.find(" g ") != -1 ]
  syms_and_address = []
  for sym in symbols:
    tokens = sym.split()
    if tokens[2] == 'F':
      address = int(tokens[0], 16)
      sym = SaneBase64(tokens[5])
      result[address] = sym
  return result

def find_nth(haystack, needle, n):
  start = haystack.find(needle)
  while start >= 0 and n > 1:
    start = haystack.find(needle, start+len(needle))
    n -= 1
  return start

def ObtainPEFunctionSymbols(training_file):
  result = {}
  filetype = subprocess.check_output(["file", "-b", training_file]).decode("utf-8")<commit_msg>Fixed the test data generation script.<commit_after>#!/usr/bin/python3

# A Python script to generate test data given ELF or PE files including their
# relevant debug info.
#
# This Python script is meant to replace the horrible bash script that is
# currently responsible for doing this.

import subprocess, os, fnmatch, codecs, random, shutil
from subprocess import Popen, PIPE, STDOUT
from operator import itemgetter
from collections import defaultdict

work_directory="/var/tmp/train2"
shutil.rmtree(work_directory)
os.mkdir(work_directory)

def FindELFTrainingFiles():
  """ Returns the list of ELF files that should be used for training. These
  ELF files need to contain objdump-able debug information.

  """
  dirnames_and_masks = [ ('./unrar.5.5.3.builds', 'unrar.x??.O?') ]

  filenames = []
  for dirname, mask in dirnames_and_masks:
    filenames = filenames + [
      dirname + "/" + filename for filename in os.listdir(dirname)
      if fnmatch.fnmatch(filename, mask) ]
  return filenames

def FindPETrainingFiles():
  """ Returns the list of PE files that should be used for training. These
  PE files need to have associated text files (with suffix .debugdump) that
  contains the output of dia2dump in the same directory. """
  return [
   "./unrar.5.5.3.builds/VS2015/unrar32/Release/UnRAR.exe",
   "./unrar.5.5.3.builds/VS2015/unrar32/Debug/UnRAR.exe",
   "./unrar.5.5.3.builds/VS2015/unrar64/Release/UnRAR.exe",
   "./unrar.5.5.3.builds/VS2015/unrar64/Debug/UnRAR.exe" ]

def ObtainFunctionSymbols(training_file, file_format):
  if file_format == "ELF":
    return ObtainELFFunctionSymbols(training_file)
  elif file_format == "PE":
    return ObtainPEFunctionSymbols(training_file)

def SaneBase64(input_string):
  """ Because Python3 attempts to win 'most idiotic language ever', encoding a
  simple string as base64 without risking to have strange newlines added is
  difficult. This functions is an insane solution: Call command line
  base64encode instead of dealing with Python. """
  encoded_string = subprocess.run(["base64", "-w0"], stdout=PIPE,
    input=bytes(input_string, encoding="utf-8")).stdout.decode("utf-8")
  return encoded_string

def ObtainELFFunctionSymbols(training_file):
  """ Runs objdump to obtain the symbols in an ELF file and then returns a
  dictionary for this file. """
  result = {}
  symbols = [ line for line in subprocess.check_output(
    [ "objdump", "-t", training_file ] ).decode("utf-8").split("\n")
      if line.find(" F .text") != -1 ]
  syms_and_address = []
  for sym in symbols:
    tokens = sym.split()
    if tokens[2] == 'F':
      address = int(tokens[0], 16)
      sym = SaneBase64(tokens[5])
      result[address] = sym
  return result

def find_nth(haystack, needle, n):
  start = haystack.find(needle)
  while start >= 0 and n > 1:
    start = haystack.find(needle, start+len(needle))
    n -= 1
  return start

def ObtainPEFunctionSymbols(training_file):
  result = {}
  filetype = subprocess.check_output(["file", "-b", training_file]).decode("utf-8")
{'fraction_non_alphanumeric': 0.1153463705667882, 'fraction_numerical': 0.022207490884985085, 'mean_word_length': 4.037562604340567, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 7, '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': '2951511', 'n_tokens_mistral': 2012, 'n_tokens_neox': 1955, 'n_words': 581}