run_bash_files.py

"""
<Program Name>
  run_bash_files.py

<Started>
  January 10, 2018.

<Last Modified>
  January 24, 2018.

<Author>
  Artiom Baloian <artiom.baloian@nyu.edu>
  Yiwen Li <liyiwen@nyu.edu>

<Copyright>
  See LICENSE for licensing information.

<Purpose>
  Find all bash scripts on Unix based OS and run them. Current module we use
  for experimenting Lind project's caging.
"""


#!/usr/bin/env python3

import os
import sys
import subprocess
import shutil
import time as timestamp
import argparse

from resource import getrusage as resource_usage, RUSAGE_SELF


# Global but not constant variables to save final result of parsing
# PERF_RESULT_FILE file.
syscall_invoked_times_dict = dict()
syscall_timing_dict = dict()
syscall_time_start = dict()
syscall_time_finish = dict()



def get_unix_time(bash_script):
  """
  Function runs bash script returns `real`, `sys` and `user` elapsed time,
  like UNIX's command `time`. You can calculate the amount of used CPU-time
  used by your function/callable by summing `user` and `sys`. `real` is just
  like the wall clock.
  Note that `sys` and `user`'s resolutions are limited by the resolution of
  the operating system's software clock (check `man 7 time` for more
  details)..

  Arguments:
  - bash_script: Bash script, note that you should provide script with
                 absolute path.
  """

  temp_unix_time_file = 'unix_time.txt'
  command = 'time -o ' + temp_unix_time_file + ' -p bash ' + bash_script
  subprocess.call(command, shell = True)

  time = {}

  if os.path.isfile(temp_unix_time_file):
    unix_time_file = open(temp_unix_time_file, 'r')
    file_data = unix_time_file.read()

    unix_time = file_data.split('\n')
    for time_str in unix_time:
      if time_str:
        temp_array = time_str.split(' ');
        key = temp_array[0]
        value = temp_array[1]

        time[key] = float(value)

    unix_time_file.close()
  else:
    print("File: ", temp_unix_time_file, " does not exist!")
    sys.exit(2)

  if os.path.exists(temp_unix_time_file):
    os.remove(temp_unix_time_file)

  return time





def collect_perf_record(bash_script):
  """
  Function runs the following command to let perf collect the system call timing
  info for running each script: (you may need "sudo")
  perf record -m 10M -e syscalls:sys_enter_* -e syscalls:sys_exit_* bash ./script_name

  Arguments:
  - bash_script: Bash script, note that you should provide script with
                 absolute path.
  """
  command = 'sudo perf record -m 10M -e syscalls:sys_enter_* -e syscalls:sys_exit_* bash ' + bash_script
  print(command)
  subprocess.call(command, shell = True)





def parse_perf_result_file(perf_result_file):
  """
  Function parses the system call time information generated by the perf_event
  tool, raw data was generated with the following command:
  perf record -m 10M -e syscalls:sys_enter_* -e syscalls:sys_exit_* bash ./script
  The raw data should look like this:

    bash 19066 [001] 329000.701687: syscalls:sys_enter_brk: brk: 0x00000000
    bash 19066 [001] 329000.701690: syscalls:sys_exit_brk: 0xb25000
    bash 19066 [001] 329000.701710: syscalls:sys_enter_access: filename: 0x7f15f5ab98c3, mode: 0x00000000
    bash 19066 [001] 329000.701716: syscalls:sys_exit_access: 0xfffffffffffffffe
  """

  # Clear dictionaries as they contain data from previous reesult.
  syscall_invoked_times_dict.clear()
  syscall_timing_dict.clear()
  syscall_time_start.clear()
  syscall_time_finish.clear()


  is_syscall_enter = -1
  for line in perf_result_file:
     if line[0] == '#':
       continue
     words = line.split(":")
     words_0 = words[0].split("] ")
     syscall_status = words[2].split("_")[1]
     if syscall_status == "enter":
       is_syscall_enter = 1
       syscall_name = words[2][10:]
     if syscall_status == "exit":
       is_syscall_enter = 0
       syscall_name = words[2][9:]

     if is_syscall_enter == 1:
       syscall_time_start[syscall_name] = float(words_0[1])
     else:
       syscall_time_finish[syscall_name] = float(words_0[1])

     if is_syscall_enter == 0:
       if syscall_name in syscall_invoked_times_dict:
         syscall_invoked_times_dict[syscall_name] = syscall_invoked_times_dict[syscall_name] + 1
         syscall_timing_dict[syscall_name] = syscall_timing_dict[syscall_name] + syscall_time_finish[syscall_name] - syscall_time_start[syscall_name]
       else:
         syscall_invoked_times_dict[syscall_name] = 1
         syscall_timing_dict[syscall_name] = syscall_time_finish[syscall_name] - syscall_time_start[syscall_name]





def get_final_result(perf_result_file, final_result_file, unix_time):
  """
  Function takes as an argument perf recored data and parses it to readable
  form. The result will save in a "final_result" directory.

  Arguments:
    - perf_result_file: Perf recor data.
    - final_result_file: In this file will be written parsed perf record.
    - unix_time:  A dictionary, which contains unix system's 'real', 'sys' and
                  'user' time.
  """

  if os.path.isfile(perf_result_file):
    perf_result_file = open(perf_result_file, 'r')
    parse_perf_result_file(perf_result_file)
    perf_result_file.close()
  else:
    print("File: ", perf_output_file, " does not exist!")
    sys.exit(2)


  final_result_file = open(final_result_file, 'w')

  final_result_file.write("real_time : " + str(unix_time["real"]) + "\n");
  final_result_file.write("user_time : " + str(unix_time["user"]) + "\n");
  final_result_file.write("sys_time : " + str(unix_time["sys"]) + "\n\n");


  syscall_total_time = 0.0
  syscall_total_invoked_times = 0
  for syscall in syscall_invoked_times_dict:
    #print(syscall + ": " + str(syscall_invoked_times_dict[syscall]) + " " + str(syscall_timing_dict[syscall]))
    syscall_total_time = syscall_total_time + syscall_timing_dict[syscall]
    syscall_total_invoked_times = syscall_total_invoked_times + syscall_invoked_times_dict[syscall]
    final_result_file.write(syscall + ": " + str(syscall_invoked_times_dict[syscall]) + " " + str(syscall_timing_dict[syscall]) + "\n")

  #print("syscall_total_time = " + str(syscall_total_time))
  final_result_file.write("\nsyscall_total_time = " + str(syscall_total_time) + "\n")

  #print("syscall_total_invoked_times = " + str(syscall_total_invoked_times))
  final_result_file.write("syscall_total_invoked_time = " + str(syscall_total_invoked_times) + "\n")

  final_result_file.close()





def find_bash_scripts_and_run(root_dir = '/'):
  """
  Function finds all bash scripts in system and runs, gets unix time,
  collects system call records, parses result of perf record inot the readable
  form, in order to see real result of bash commands.
  Also, it creates "perf_result" and "final_result" where will be stored result
  of perf record and parsed version of perl record.

  Every single script will have it's own perf record in a "perf_result" and
  parsed perf record in a "final_result" directories respectively.
  In "perf_result" directory filenames format will be as follow:
  <script_name>_perf_output.txt
  In "final_result" directory filenames format will be as follow:
  <script_name>_perf__final_output.txt

  Arguments:
    - root_dir: Root directory from where it will start looking for bash scripts.
                By default it is set root of OS.

  Returns:
   - A number od bash script run by function. 
  """

  # Create a 'perf_result' directory where is going to be stored result 
  # of perf record data in plain text.
  perf_directory = "perf_result"

  # Remove old data if it exists
  if os.path.exists(perf_directory):
    shutil.rmtree(perf_directory)

  try:
      os.makedirs(perf_directory)
  except OSError as e:
    if e.errno != errno.EEXIST:
      raise


  # Create a 'final_result' directory where is going to be stored result 
  # of final data after parsing perf record.
  final_directory = "final_result"

  # Remove old data if it exists
  if os.path.exists(final_directory):
    shutil.rmtree(final_directory)

  try:
      os.makedirs(final_directory)
  except OSError as e:
    if e.errno != errno.EEXIST:
      raise


  # Keep a number of bash script files.
  count = 0;
  for sub_dir, dirs, files in os.walk(root_dir):

    # Parse and skip 'home' directory, because we are noot looking for bash
    # script inside of 'home' directory.
    # Uncomment these two lines if you want to skip 'home' directory.
    #is_home_dir = sub_dir[0:6]
    #if is_home_dir != "/home/":

    for file in files:
      # TODO There are bash scripts, which are not ending with '.sh'
      # Idea is to open file and look bash commands in order to verify it.
      if file.endswith('.sh'):
        count = count + 1;
        script_absolute_path = sub_dir + '/' + file
        script_file_name = (str(file)).split(".")[0]

        unix_time = get_unix_time(script_absolute_path)

        collect_perf_record(script_absolute_path)

        # When you are done with perf record, there should be a perf.data file
        # created by perf, which contains the data we wanted. 
        # Use 'sudo perf script > perf_result_output.txt' command to write the data
        # to a plain txt file.
        perf_result_file = perf_directory + '/' + script_file_name + "_perf_output.txt"
        command = 'sudo perf script > ' + perf_result_file
        subprocess.call(command, shell = True);

        #TODO It is supposed to call a function here, which removes the unuseful
        # info at the beginning of file. For now it is OK.

        final_result_file = final_directory + '/' + script_file_name + "_final_output.txt"
        get_final_result(perf_result_file, final_result_file, unix_time)

  return count





def main():

  parser = argparse.ArgumentParser()
  parser.add_argument("--root_dir", required = False, help = "root directory", type = str, default = '/')

  argv = parser.parse_args(sys.argv[1:])
  root_directory = argv.root_dir

  count = find_bash_scripts_and_run(root_directory)



if __name__ == "__main__":
  main()