# -*- coding: binary -*-

module Rex
  module Exploitation
    #
    # VBScript obfuscation library
    #
    class VBSObfuscate
      # The VBScript code that this obfuscator will transform
      attr_accessor :code

      # Saves +code+ for later obfuscation with #obfuscate!
      #
      # @param code [#to_s] the code to obfuscate
      # @param opts [Hash] an options hash
      def initialize(code = nil, _opts = {})
        self.code = code
      end

      # @return [String] the (possibly obfuscated) code
      def to_s
        @code
      end

      # Append +str+ to the (possibly obfuscated) code
      def <<(str)
        @code << str
      end

      # Obfuscate VBScript code.
      #
      # @option iterations            [Integer] number of times to run the obfuscator on this code (1)
      # @option normalize_whitespace  [Boolean] normalize line endings and strip leading/trailing whitespace from each line (true)
      # @option dynamic_execution     [Boolean] dynamically execute obfuscated code with Execute (true)
      #
      # @return [self]
      def obfuscate!(iterations: 1, normalize_whitespace: true, dynamic_execution: true)
        raise(ArgumentError, 'code must be present') if @code.nil?
        raise(ArgumentError, 'iterations must be a positive integer') unless iterations.integer? && iterations.positive?

        obfuscated = @code.dup

        iterations.times do
          # Normalize line endings and strip leading/trailing whitespace
          if normalize_whitespace
            obfuscated.gsub!(/\r\n/, "\n")
            obfuscated = obfuscated.lines.map(&:strip).reject(&:empty?).join("\n")
          end

          # Convert all VBScript to a string to be dynamically executed with Execute()
          if dynamic_execution
            obfuscated = 'Execute ' + vbscript_string_for_execute(obfuscated)
          end

          # Obfuscate strings
          obfuscated = chunk_vbscript_strings(obfuscated)
          obfuscated.gsub!(/"((?:[^"]|"")*)"/) do
            raw = ::Regexp.last_match(1).gsub('""', '"')
            raw.chars.map { |c| "chr(#{generate_number_expression(c.ord)})" }.join('&')
          end

          # Obfuscate integers
          obfuscated.gsub!(/\b\d+\b/) do |num|
            generate_number_expression(num.to_i)
          end
        end

        @code = obfuscated

        self
      end

      private

      # Converts all VBScript in +vbscript+ to a string for dynamic execution
      # with Execute().
      #
      # @param vbscript [String] VBScript code
      #
      # @return [String] obfuscated VBScript code for use with Execute()
      def vbscript_string_for_execute(vbscript)
        lines = vbscript.lines.map(&:chomp).map do |line|
          escaped_line = line.gsub('"', '""')
          "\"#{escaped_line}\""
        end
        lines.join('&vbCrLf&')
      end

      # Returns a random math expression evaluating to input +int+
      #
      # @param int [Integer] input integer
      #
      # @return [String] math expression evaluating to input +int+
      def generate_number_expression(int)
        case rand(4)
        when 0 # Sum
          a = rand(0..int)
          b = int - a
          "(#{a}+#{b})"
        when 1 # Difference
          r1 = int + rand(1..10)
          r2 = r1 - int
          "(#{r1}-#{r2})"
        when 2 # Product (only if divisible)
          divisors = (1..int).select { |d| (int % d).zero? }
          if divisors.size > 1
            d = divisors.sample
            "(#{d}*#{int / d})"
          else
            "(#{int}+0)"
          end
        when 3 # Quotient
          r2 = rand(1..10)
          r1 = int * r2
          "(#{r1}/#{r2})"
        end
      end

      # Return VBScript code with all strings split into chunks and concatenated
      #
      # @param vbscript [String] VBScript code
      #
      # @return [String] VBScript code with chunked strings
      def chunk_vbscript_strings(vbscript)
        vbscript.gsub(/"([^"]+)"/) do
          original = Regexp.last_match(1)
          chunks = []

          i = 0
          while i < original.length
            chunk_size = rand(1..5)
            chunks << "\"#{original[i, chunk_size]}\""
            i += chunk_size
          end

          chunks.join('&')
        end
      end
    end
  end
end
