#!/usr/bin/ruby
#
# stoat - LLVM Based Static Analysis Tool
# Copyright (C) 2015 Mark McCurry
#
# This file is part of stoat.
#
# stoat is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# stoat is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with stoat.  If not, see <http://www.gnu.org/licenses/>.
#
require 'optparse'
require 'ostruct'
require 'set'
require 'yaml'
require 'pp'
require 'csv'
#require 'ruby-prof'
#RubyProf.start

CurrentDir = File.dirname(__FILE__)

#Load Dependencies
def load_dep(x)
    load "#{CurrentDir}/lib/#{x}"
rescue LoadError
    load "#{CurrentDir}/../share/stoat/#{x}"
end

load_dep "callgraph.rb"
load_dep "deductions.rb"
load_dep "graph.rb"
load_dep "load-callgraph.rb"

################################################################################
#                            Parse Options                                     #
################################################################################

options = OpenStruct.new
options.whitelist = [CurrentDir+"/data/whitelist.txt",
                     CurrentDir+"/../share/stoat/whitelist.txt"]
options.blacklist = [CurrentDir+"/data/blacklist.txt",
                     CurrentDir+"/../share/stoat/blacklist.txt"]
options.suppression = []
options.unmangled = OpenStruct.new
options.unmangled.whitelist = []
options.unmangled.blacklist = []
options.root = "./"
options.lib  = "libstoat.so"
options.dir  = []
options.graphfile = nil
options.recursive = false
options.dump_file = nil
options.minimal_graph = false
options.shorten = false
options.colorize = false

OptionParser.new do |opts|
      opts.banner = "Usage: stoat [options] [FILES]"

      opts.on("-w", "--whitelist FILE",
              "Define a Whitelist File") do |list|
          options.whitelist << list
      end

      opts.on("-b", "--blacklist FILE",
              "Define a Blacklist File") do |list|
          options.blacklist << list
      end

      opts.on("-s", "--suppression FILE",
              "Define a Suppression File") do |list|
          options.suppression << list
      end

      opts.on("-r", "--recursive DIR",
              "Enable Recursive Search Mode (this makes FILES optional)") do |dir|
          options.recursive = true
          options.root = dir
      end

      opts.on("-l", "--llvm-passes LIB",
              "The Library Containing The Needed LLVM Passes") do |lib|
          options.lib = lib
      end

      opts.on("-g", "--graph-view FILE.png",
              "The Graph View Output File Name") do |file|
          options.graphfile = file
      end

      opts.on("-G", "--graph-minimal-view FILE.png",
              "The Minimal Graph View Output File Name") do |file|
          options.graphfile = file
          options.minimal_graph = true
      end

      opts.on("-S", "--shorten-names",
              "Omit Namespaces and Template Arguments in Graph") do
          options.shorten = true
      end

      opts.on("-d", "--dump FILE.txt",
              "Dump Information Extracted From LLVM IR") do |file|
          options.dump_file = file
      end

      opts.on("-c", "--color",
              "Colorize Output") do
          options.colorize = true
      end
end.parse!

################################################################################
#                        Load Optional Dependencies                            #
################################################################################
if(options.graphfile)
    begin
        require 'graphviz'
        test = GraphViz
    rescue LoadError
        puts "ruby-graphviz gem could not be loaded"
        puts "please install it with 'gem install ruby-graphviz'"
        exit
    rescue NameError
        puts "ruby-graphviz gem could not be loaded"
        puts "to fix this"
        puts "1) remove conflicting 'graphviz' gem with 'gem uninstall graphviz'"
        puts "2) install it with 'gem install ruby-graphviz'"
        exit
    end
end

if(options.colorize)
    require 'colorize'
else
    class String
        def bold
            self
        end
        def colorize(color)
            self
        end
    end
end

################################################################################
#                            Load Data Files                                   #
################################################################################

def read_file(file)
    return File.read(file).split("\n") \
        .select{|x| /^\s*(#.*)?$/ !~ x} \
        .map{|x| CSV.parse(x, col_sep: ' ')}.flatten(1)
end

#White List/Black List Expansion
white_list = []
tmp = 0
options.whitelist.each do |x|
    begin
        tmp = tmp + 1
        white_list.concat read_file(x).flatten(1)
    rescue Errno::ENOENT => err
        if(tmp > 2) #Non default
            $stderr.puts("Warning: Unknown whitelist file '#{x}'".colorize(:yellow))
        end
    end
end

black_list = []
tmp = 0
options.blacklist.each do |x|
    begin
        tmp = tmp + 1
        black_list.concat read_file(x).flatten(1)
    rescue Errno::ENOENT => err
        if(tmp > 2) #Non default
            $stderr.puts("Warning: Unknown blacklist file '#{x}'".colorize(:yellow))
        end
    end
end

suppression_list = []
begin
    tlist = []
    options.suppression.each do |x|
        begin
            tlist.concat read_file(x)
        rescue Errno::ENOENT => err
            $stderr.puts("Warning: Unknown suppression file '#{x}'".colorize(:yellow))
        end
    end
    tlist.each do |x|
        if(x.length == 3 and x[1] == '==>')
            tmp = x[0], x[2]
            suppression_list << tmp
        else
            $stderr.puts("Error: Possible formatting error in suppression file".colorize(:red))
        end
    end
end


################################################################################
#                      Get Parameters For LLVM Passes                          #
################################################################################
files = []
if(!ARGV.empty?)
    files.concat ARGV
end
if(options.recursive)
    rfiles = []
    rfiles.concat `find #{options.root} -type f | grep -e "\\.bc$"`.split
    rfiles.concat `find #{options.root} -type f | grep -e "\\.o$"`.split

    rfiles.each do |f|
        if(/LLVM/.match `file #{f}`)
            files << f
        end
    end
end

if(files.empty?)
    $stderr.puts("Error: There Are No Files To Process".colorize(:red))
    exit 1
end

#Identify opt binary name
opt = nil
["opt", "opt-4.0", "opt-3.9", "opt-3-7", "opt-3.6", "opt-3.5", "opt-3.4", "opt-3.3"].each do|x|
    if(!`which #{x} 2> /dev/null`.empty?)
        opt = x
        break
    end
end

if(!opt)
    $stderr.puts("Error: Could not find a usable version of 'opt'".colorize(:red))
    exit 1
end

library_test = `#{opt} -load #{options.lib} < #{__FILE__} 2>&1`
if(/Error opening/.match library_test)
    $stderr.puts("Error: Could not find a usable libstoat.so".colorize(:red))
    exit 1
end

################################################################################
#                           Initialize Callgraph                               #
################################################################################

(callgraph, classes, vtable, rtosc) = load_callgraph(opt, options.lib, files)

#Add Anything That's On the function_props list
Deductions::add_attr_safety(callgraph)

#Add Rtosc information
Deductions::add_rtosc_safety(callgraph, rtosc)

#Add Any Known Virtual Calls
callgraph.add_virtual_methods vtable

#Add Calls Down the hierarchy
callgraph.add_subclass_calls(classes,vtable)

#Add C++ABI Destructor/Constructor Chaining
callgraph.add_constructor_chains
callgraph.add_destructor_chains

#Rebuild Demangled Name Cache
callgraph.update_demangled_cache rtosc

################################################################################
#                              Apply Datafiles                                 #
################################################################################

#Add WhiteList Information
callgraph.apply_whitelist white_list

#Add BlackList Information
callgraph.apply_blacklist black_list

#Suppress any bit of the callgraph which is noted in the suppression file
callgraph.apply_suppression suppression_list

#Add no source stuff
Deductions::add_no_body_safety callgraph

#puts callgraph.full_to_s

################################################################################
#                              Save Datafiles                                  #
################################################################################

#Dump Known information to a file
if(options.dump_file)
    File.open(options.dump_file, "w") do |file|
        file.puts callgraph.full_to_s
    end
end

################################################################################
#                             Perform Deductions                               #
################################################################################

deductions = Deductions::setup(callgraph)
Deductions::deduce_safe(deductions, callgraph)
error_count = Deductions::dump_errors(deductions, callgraph)
Deductions::print_stats deductions

################################################################################
#                               Save Graph                                     #
################################################################################

if options.graphfile
    GraphRender::to_graph(deductions, callgraph, options.graphfile,
                          options.minimal_graph, options.shorten)
end
puts "return code is '#{(error_count != 0) ? -1 : 0}'"

#result = RubyProf.stop
#printer = RubyProf::GraphPrinter.new(result)
#printer.print(STDOUT, {})

exit(-1) unless error_count == 0

