require 'pentest/checkers/base_checker'
require 'pentest/payload'
require 'pentest/sql_proxy'
require 'term/ansicolor'
require 'pairwise'
require 'arproxy'
require 'callsite'
require 'gda'

class Pentest::SqliChecker < Pentest::BaseChecker
  Pentest::Checkers.add self

  @description = "Checks for SQL injections"

  SQLI_PAYLOADS = File.read(File.expand_path('../fuzzers/sqli.txt', File.dirname(__FILE__)), encoding: 'utf-8').lines.map(&:strip).select {|l| l.size > 5 && l =~ /\W/}
  CRACKER_PAYLOAD = %q(<>"'%;()&+)

  def initialize(endpoint, params)
    super(endpoint, params)
    @queries = []
    @parser = GDA::SQL::Parser.new
  end

  def generate_preattack_payloads(params, seeds, injection_point)
    values_list = if params.size - 1 <= 0
      [[]]
    elsif params.size - 1 == 1
      seeds.map {|s| [s]}
    else
      Pairwise.combinations(*([seeds] * (params.size - 1)))
    end

    values_list.map do |values|
      values.insert(injection_point, CRACKER_PAYLOAD)

      Pentest::Payload.new(
        params: params,
        route: @route,
        values: values,
        injection_point: injection_point,
        injection: CRACKER_PAYLOAD,
      )
    end.take(50)
  end

  def handle_query(sql)
    @queries << [sql, caller]
  end

  def attack(param, injection_point, ingredients)
    preattack_payloads = generate_preattack_payloads(@params, ingredients, injection_point)

    errors = []

    penetrated_payload = nil
    preattack_payloads.shuffle.each do |payload|
      request, response, err = dispatch(payload)
      status = get_status(err) || response.status

      Pentest::Logger.put_progress (status / 100).to_s

      errors << normalize_error(err, payload)

      if ::ActiveRecord::StatementInvalid === err
        payload.penetration_type = 'SQL Injection Vulnerability'
        payload.penetration_confidence = :preattack
        penetrated_payload = payload
        break
      end
    end

    # preattack not succeeded. skipping
    return [nil, errors] if penetrated_payload.nil?

    # attack
    attack_payloads = generate_attack_payloads(@params, penetrated_payload.values, injection_point)

    Pentest::SqlProxy.enable!(self.method(:handle_query))

    attack_payloads.shuffle.each do |payload|
      request, response, err = dispatch(payload)
      status = get_status(err) || response.status

      Pentest::Logger.put_progress (status / 100).to_s

      errors << normalize_error(err, payload)

      check_attack_result(payload, err);

      if payload.penetration_confidence == :attack
        penetrated_payload = payload
        break
      end
    end

    Pentest::SqlProxy.disable!(self.method(:handle_query))

    [penetrated_payload, errors]
  end

  def generate_attack_payloads(params, values, injection_point)
    SQLI_PAYLOADS.map do |injection|
      new_values = values.dup
      new_values[injection_point] = injection

      Pentest::Payload.new(
        params: params,
        route: @route,
        values: new_values,
        injection_point: injection_point,
        injection: injection,
      )
    end
  end

  def check_attack_result(payload, err)
    @queries.each do |query, trace|
      begin
        stmt = @parser.parse(query.tr('`?', '"0'))
      rescue RuntimeError => e
        next
      end

      if query.include?(payload.injection) && !(GDA::Nodes::Unknown === stmt.ast)
        tokens = extract_values(stmt.ast)
        if tokens.all? {|token| !token.force_encoding("UTF-8").include? payload.injection}
          callsites = Callsite.parse(trace)
          project_call = callsites.find do |callsite|
            callsite.filename.starts_with?(@app_path)
          end
          unless project_call.nil?
            line = File.read(project_call.filename).lines[project_call.line - 1].rstrip
          end

          payload.penetration_type = 'SQL Injection Vulnerability'
          payload.penetration_confidence = :attack
          payload.penetration_message = [
            *(line.nil? ? [] : [
              "#{project_call.filename}:#{project_call.line}",
              "> #{line}",
            ]),
            "Issued query: #{query.gsub(payload.injection, Term::ANSIColor.red(payload.injection))}",
          ].join("\n")
          break
        end
      end
    end
  end

  def extract_values(stmt)
    ret = []

    if stmt.nil?
      # nop
    elsif stmt.is_a?(String)
      ret << stmt
    elsif stmt.is_a?(Integer)
      ret << stmt.to_s
    elsif stmt.is_a?(Array)
      stmt.each do |s|
        ret += extract_values(s)
      end
    elsif GDA::Nodes::Table === stmt
      ret += extract_values(stmt.table_name)
    elsif GDA::Nodes::Field === stmt
      ret += extract_values(stmt.field_name)
    elsif GDA::Nodes::Expr === stmt
      ret += extract_values(stmt.func)
      ret += extract_values(stmt.cond)
      ret += extract_values(stmt.select)
      ret += extract_values(stmt.case_s)
      ret += extract_values(stmt.param_spec)
      ret += extract_values(stmt.cast_as)
      ret << stmt.value
    elsif GDA::Nodes::Select === stmt
      ret += extract_values(stmt.distinct_expr)
      ret += extract_values(stmt.expr_list)
      ret += extract_values(stmt.from)
      ret += extract_values(stmt.where_cond)
      ret += extract_values(stmt.group_by)
      ret += extract_values(stmt.having_cond)
      ret += extract_values(stmt.order_by)
      ret += extract_values(stmt.limit_count)
      ret += extract_values(stmt.limit_offset)
    elsif GDA::Nodes::SelectField === stmt
      ret += extract_values(stmt.expr)
      ret += extract_values(stmt.field_name)
      ret += extract_values(stmt.table_name)
      ret += extract_values(stmt.as)
    elsif GDA::Nodes::From === stmt
      ret += extract_values(stmt.targets)
      ret += extract_values(stmt.joins)
    elsif GDA::Nodes::Operation === stmt
      ret += extract_values(stmt.operands)
    elsif GDA::Nodes::Target === stmt
      ret += extract_values(stmt.expr)
      ret += extract_values(stmt.table_name)
      ret += extract_values(stmt.as)
    elsif GDA::Nodes::Function === stmt
      ret += extract_values(stmt.args_list)
      ret += extract_values(stmt.function_name)
    elsif GDA::Nodes::Order === stmt
      ret += extract_values(stmt.expr)
      ret += extract_values(stmt.collation_name)
    elsif GDA::Nodes::Insert === stmt
      ret += extract_values(stmt.table)
      ret += extract_values(stmt.fields_list)
      ret += extract_values(stmt.expr_list)
      ret += extract_values(stmt.cond)
      ret += extract_values(stmt.conflict)
    elsif GDA::Nodes::Delete === stmt
      ret += extract_values(stmt.table)
      ret += extract_values(stmt.cond)
    elsif GDA::Nodes::Join === stmt
      ret += extract_values(stmt.expr)
      ret += extract_values(stmt.use)
      ret += extract_values(stmt.position)
    elsif GDA::Nodes::Compound === stmt
      ret += extract_values(stmt.compound_type)
    elsif GDA::Nodes::Unknown === stmt
      ret += extract_values(stmt.expressions)
    end

    ret
  end
end