#
#   Copyright (C) Stonesoft Corporation 2010 - 2012.
#   All rights reserved.
#
#   The StoneGate software, manuals, and technical
#   literature may not be reproduced in any form or
#   by any means except by permission in writing from
#   Stonesoft Corporation.
#

# Predator4 webGUI

require 'rubygems'
require 'sinatra'
require 'sinatra/base'
require 'erb'
require 'json'
require "#{$webgui_root_webgui}/wg_helpers"
require "#{$webgui_root_webgui}/wg_evasion_parser"
require "#{$webgui_root_webgui}/wg_pdf_creator"

module P4WebGUI
  @@web_log = []
  @@mutex = Mutex.new     # Hold when doing anything that changes our state
  @@exploit_inst = nil    # Current exploit instance
  @@log_mutex = Mutex.new # Hold while changing log offset table
  @@log_offsets = {}
      
  class HTTPHandler < Sinatra::Base
    set :logging, true
    set :sessions, true
    set :public_folder, "#{$webgui_root_webgui}/public"
    set :views, "#{$webgui_root_webgui}/public/templates"
    set :environment, :development
    set :server, %w[ thin webrick ]
    set :bind, "0.0.0.0"
    mime_type :pcap, 'application/pcap'

    helpers do
      def base_url
        @base_url ||= "#{request.env['rack.url_scheme']}://#{request.env['HTTP_HOST']}"
      end
    end

    get '/' do
      @config = DataStore::get_config
      @p4 = DataStore::get_predator4
      erb :index
    end

    post '/exploit/run' do
      P4WebGUI::exploit_run(params)
      ""
    end

    get '/exploit/stop' do
      P4WebGUI::exploit_stop
      ""
    end
    
    post '/exploit/report' do
      body = JSON.parse(request.body.read)
      pdf_creator = P4WebGUI::PDFCreator.new settings.public_folder
      filename = pdf_creator.create_pdf body['log'], body['duration'], body['breaches']
      {"url" => "#{base_url}/reports/#{filename}"}.to_json
    end
    
    get '/exploit_outcome' do
      P4WebGUI::exploit_outcome_for(params[:cmd])      
    end

    get '/log' do
      P4WebGUI::generate_log_html(params)
    end

    get '/shell/read' do
      P4WebGUI::shell_read(params)
    end

    get '/shell/write' do
      P4WebGUI::shell_write(params)
    end

    get '/shell/close' do
      P4WebGUI::shell_close
    end

    get '/shell/reader_html' do
      erb :shell, :layout => false
    end
  end
  
  def self.generate_log_html(params)
    session_id = params["session_id"]
    if Helpers.blank?(session_id)
      log_error("generate_log_html - params does not include session_id!")
      return ""
    end
    
    @@log_mutex.synchronize do      
      @@log_offsets[session_id] ||= 0

      curr_offset = @@log_offsets[session_id]
      if curr_offset > @@web_log.size
        curr_offset = @@log_offsets[session_id] = @@web_log.size
      end

      output_str = @@web_log[curr_offset..-1]

      @@log_offsets[session_id] += output_str.size
      output_str
    end
  end
  
  def self.exploit_outcome_for(cmd)
    Predator4Module::get_evasion_layers_from_attack(cmd)
  end

  def self.shell_read(params)
    @@command_shell_mutex.synchronize do
      if @@command_shell_inst.nil?
        "No shell instance!"
      else
        @@command_shell_inst.get_content_html(params)
      end
    end
  end

  def self.shell_write(params)
    @@command_shell_mutex.synchronize do
      if @@command_shell_inst.nil?
        "No shell instance!"
      else
        @@command_shell_inst.write_content(params)
      end
    end
  end

  def self.shell_close
    @@command_shell_mutex.synchronize do
      unless @@command_shell_inst.nil?
        @@command_shell_inst.close
        @@command_shell_inst = nil
      end
      ""
    end
  end
  
  def self.exploit_run(params)
    @@web_log = []
    @@log_offsets = {}

    unless params["cmd"].nil?
      cmd = params["cmd"].strip
      unless cmd.match(/shell_tcp/)
        cmd = "#{cmd} --shell_tcp --record=#{$webgui_tmp_pcap}"
      end
      log_path= "log/writable/"
    else
      cfg_src_port = cfg_dst_port = cfg_dst_mask = cfg_gw_if = nil
      host_cfg = network_cfg = mode_cfg_str = tool_name = nil
      network_cfg_str = ""

      p4 = DataStore::get_predator4

      cfg_iface = params["netconfig_iface"]
      cfg_src_ip = params["netconfig_src_ip"]
      cfg_src_mask = nil
      cfg_dst_ip = params["netconfig_dst_ip"]
      cfg_gw = params["netconfig_gateway"]

      if cfg_src_ip =~ /^(.+)\/(\d+)\s*$/
        cfg_src_ip = $1
        cfg_src_mask = $2.to_i
      end

      if cfg_src_ip.match(/-/)
        range = Helpers::IPRange.new(cfg_src_ip)
        workers = range.workers
        cfg_src_ip = range.src_ip
      end

      if cfg_src_mask.nil? and not params["netconfig_src_mask"].nil? and params["netconfig_src_mask"].to_s.strip.size != 0
        cfg_src_ip = "#{cfg_src_ip}/#{params["netconfig_src_mask"]}"
      elsif not cfg_src_mask.nil?
        cfg_src_ip = "#{cfg_src_ip}/#{cfg_src_mask}"
      else
        cfg_src_ip = "#{cfg_src_ip}"
      end

      cfg_attack = params["attack"]
      attack_cfg_str = ""
      if cfg_attack.nil?
        log_web_error( "Invalid exploit URL - attack missing" )
        return
      end

      attack_info = p4.attacks_by_name[cfg_attack]
      if attack_info.nil?
        log_web_error("Invalid exploit URL - attack \"#{cfg_attack}\" not found")
        return
      end
      attack_cfg_str << "--attack=#{cfg_attack}"

      mode = params["mode"]
      if mode == "mongbat"
        tool_name = "ruby mongbat.rb --uid=#{$uid}"

        network_cfg_str << "--iface=#{cfg_iface} " unless cfg_iface.nil?
        network_cfg_str << "--attacker=#{cfg_src_ip} " unless cfg_src_ip.nil?
        network_cfg_str << "--victim=#{cfg_dst_ip} " unless cfg_dst_ip.nil?
        network_cfg_str << "--gw=#{cfg_gw} " unless Helpers.blank?(cfg_gw)
        network_cfg_str << "--mask=#{cfg_src_mask} " unless Helpers.blank?(cfg_src_mask) || Helpers.blank?(cfg_gw)
        network_cfg_str.strip!

        mongbat_cfg_str = ""
        if( params["mongbat_mode"].nil? || params["mongbat_time"].nil? ||
            params["mongbat_workers"].nil? || params["mongbat_minev"].nil? ||
            params["mongbat_maxev"].nil? )

          log_web_error( "Mongbat configuration entries missing!" )
          return
        end

        params["mongbat_minev"] = "1" if params["mongbat_maxev"] == "1"

        mongbat_cfg_str = "--mode=#{params['mongbat_mode']} --time=#{params['mongbat_time']} --workers=#{workers || params['mongbat_workers']}" +
                          " --min_evasions=#{params['mongbat_minev']} --max_evasions=#{params['mongbat_maxev']} --passthrough"
        mode_cfg_str = mongbat_cfg_str
        
        unless params["attack_payload"].nil?
          exploit_payload_str = "--payload=#{params["attack_payload"]} --check_victim=false"
        end
      elsif mode == "evasions"
        tool_name = p4.p4_bin_path

        network_cfg_str << "--if=#{cfg_iface} " unless cfg_iface.nil?
        network_cfg_str << "--src_ip=#{cfg_src_ip} " unless cfg_src_ip.nil?
        network_cfg_str << "--src_port=#{cfg_src_port} " unless cfg_src_port.nil?
        network_cfg_str << "--src_prefix=#{cfg_src_mask} " unless Helpers.blank?(cfg_src_mask) || Helpers.blank?(cfg_gw)
        network_cfg_str << "--dst_ip=#{cfg_dst_ip} " unless cfg_dst_ip.nil?
        network_cfg_str << "--dst_port=#{cfg_dst_port} " unless cfg_dst_port.nil?
        network_cfg_str << "--dst_mask=#{cfg_dst_mask} " unless cfg_dst_mask.nil?
        network_cfg_str << "--gw=#{cfg_gw} " unless Helpers.blank?(cfg_gw)
        network_cfg_str.strip!

        mode_cfg_str = EvasionParser.parse(params['evasions'])
        mode_cfg_str += "--obfuscate "
      end

      exploit_extra_str = ""
      unless params["attack_option"].nil?
        exploit_extra_str = "--extra=#{params["attack_option"]}=true"
      end

      other_opt_cfg_str = ""
      begin
        other_opt_cfg_str += Helpers.add_to_option_string("autoclose", params["auto_close_shell"])
        other_opt_cfg_str += Helpers.add_to_option_string("clean", params["clean"])
      rescue => e
        log_web_error e.message
        return
      end

      if mode == "evasions"
        other_opt_cfg_str << "--shell_tcp --record=#{$webgui_tmp_pcap} "
      end

      other_opt_cfg_str.strip!

      tail_opt_cfg_str = "--verifydelay=#{$p4_verifydelay} "

      cmd = "#{tool_name} #{attack_cfg_str} #{exploit_payload_str || ''} #{network_cfg_str} #{other_opt_cfg_str} #{mode_cfg_str} #{exploit_extra_str} #{tail_opt_cfg_str}"
      cmd.strip!
      
      log_path= "log/writable/"
    end
    
    @@mutex.synchronize do
      unless @@exploit_inst.nil?
        log_web_error("Already running an exploit!")
        return
      end

      # Generate an unique ID for this attack run - current time should be unique enough as we are in a mutex
      # This is used for the packet capture and log file
      unique_id = Time.new.strftime( "%C%y%m%d_%k%M%S_%3N") # <year><month><day>_<hour><minute><second>_<milliseconds>
      register_exploit( P4Exploit.new( cmd, log_path, unique_id ) )
    end
  end

  def self.exploit_stop
    @@mutex.synchronize do
      if @@exploit_inst.nil?
        log_web_info( "No exploit to stop" )
      else
        @@exploit_inst.close
        log_web_info( "Exploit stopped" )
        @@exploit_inst = nil
      end
    end
  end

  def self.escape_log_web(str)
    escaped = str.chomp
    escaped.gsub!( "<", "&lt;" )
    escaped.gsub!( ">", "&gt;" )
    escaped.gsub!( /(\r\n|\r|\n)/, "<br>" )
    escaped
  end

  def self.log_web(str, escape = true)
    use_mutex = true

    begin
      @@mutex.lock
    rescue ThreadError
      use_mutex = false
    end

    if escape
      final_str = escape_log_web(str)
    else
      final_str = str
    end

    # Make http://<ip address>:<port>/ -strings into links
    if final_str =~ /(http:\/\/\d+\.\d+\.\d+\.\d+(:(\d+))?\/([^\x20\x09])*)/
      url = $1
      final_str.gsub!( url, "<a class=\"log_link\" href=\"#{url}\">#{url}</a>" )
    end
    
    @@web_log.push final_str
    @@mutex.unlock if use_mutex
  end

  def self.log_web_info(str, escape = true)
    use_mutex = true

    begin
      @@mutex.lock
    rescue ThreadError
      use_mutex = false
    end

    if escape
      final_str = escape_log_web(str) + "<br>"
    else
      final_str = str
    end
    
    @@web_log.push "Info: #{final_str}"
    @@mutex.unlock if use_mutex
  end

  def self.log_web_error(str, escape = true)
    use_mutex = true
   
    begin
      @@mutex.lock
    rescue ThreadError
      use_mutex = false
    end

    if escape
      final_str = escape_log_web(str) + "<br>"
    else
      final_str = str
    end
    
    @@web_log.push "Error: #{final_str}"
    @@mutex.unlock if use_mutex
  end

  # register_exploit does not take @@mutex as it is always called from within run_exploit
  def self.register_exploit(expl_inst)
    if not @@exploit_inst.nil?
      raise "Exploit not nil!"	# Ok to raise as this is checked in run_exploit
    end
		
	@@exploit_inst = expl_inst
  end
  
  def self.unregister_exploit(expl_inst)
    use_mutex = true

    begin
      @@mutex.lock
    rescue ThreadError => e
      use_mutex = false
    end

    if @@exploit_inst.nil?
    elsif @@exploit_inst != expl_inst
    else
      @@exploit_inst = nil
    end
    
    @@mutex.unlock if use_mutex    
  end
  
  def self.page_title
    $is_evader ? "Evader" : "Predator4"
  end

  log_web_info("WebGUI started on #{Time.new}")
end
