#!/usr/bin/env python
import sys
import os
import re
import xml.etree.ElementTree as ET
import argparse
#Self-defined functions
import parse

cur_path = os.path.abspath('{0}/..'.format(os.path.realpath(__file__)))

def check_extension(file_name):
  try:
    pat = '.*\.[ch]$'
    if re.match(pat, file_name):
      return True
    else:
      return False
  except Exception as e:
    print 'Extention error: {0}'.format(e)
    sys.exit()

def arg_parser():
  try:
    parser = argparse.ArgumentParser(description='Buffer Overflow Attack Detector.')
    parser.add_argument('-i', '--input', nargs='+', help='Give the input file name')
    parser.add_argument('-o', '--output', help='Specify the name of output report.')
    parser.add_argument('-a', '--append', action='store_true', default='False', help='Report appends to file.')
    args = parser.parse_args()
    return args
  except Exception as e:
    print 'argparse error: {0}'.format(e)
    sys.exit()
  
def main():

  report_name = '{0}/Report'.format(cur_path)
  illegal_file = ''

  args = arg_parser()
  if not args.input:
    print 'ERROR: Lack of input file.'
    sys.exit(-1)
  if args.output:
    report_name = '{0}/{1}'.format(cur_path, args.output)

  try:
    if args.append:
      fw = open(report_name, 'a')
    else:
      fw = open(report_name, 'w')
  except IOError as e:
    print e
    sys.exit(-1)
  
  for i in range(len(args.input)):
    input_name = os.path.abspath('{0}/{1}'.format(cur_path, args.input[i]))
    if check_extension(input_name):
      try:
        fr = open(input_name, 'r')
        in_str = 'File "{0}"\n'.format(input_name)
        print_flag = False
      except IOError as e:
        print e
        sys.exit(-1)
      try:
        #Read dangerous function list
        func_list = parse.parse_dangerous_func()
        line_count = 1
        #Read source file and try to find dangerous functions
        for line in fr.readlines():
          line = line.strip()
          for func in func_list:
            if parse.find_match(func['reg'], line):
              print_flag = True
              in_str = in_str + 'line {0} with dangerous function {1}: {2}\n'.format(line_count, func['name'], line)
          line_count = line_count + 1
      except Exception as e:
        print e
        sys.exit(-1)
      try:
        if print_flag:
          in_str = in_str + '------------------\n'
          fw.write(in_str)
        fr.close()
      except IOError as e:
        print e
        sys.exit(-1)
    else:
      illegal_file = '{0} {1}'.format(illegal_file, input_name)
  
  try:
    fw.close()
    print 'Report is stored at {0}'.format(os.path.abspath(report_name))
    if not (illegal_file is ''):
      print 'Illegal input file: {0}'.format(illegal_file)
  except IOError as e:
    print e
    sys.exit(-1)

if __name__ == "__main__":
  main()
