#
#   Copyright 2012 Wade Alcorn wade@bindshell.net
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
module BeEF
module Module 

    # Checks to see if module key is in configuration
    # @param [String] mod module key
    # @return [Boolean] if the module key exists in BeEF's configuration
    def self.is_present(mod)
        return BeEF::Core::Configuration.instance.get('beef.module').has_key?(mod.to_s)
    end

    # Checks to see if module is enabled in configuration
    # @param [String] mod module key
    # @return [Boolean] if the module key is enabled in BeEF's configuration
    def self.is_enabled(mod)
        return (self.is_present(mod) and BeEF::Core::Configuration.instance.get('beef.module.'+mod.to_s+'.enable') == true)
    end

    # Checks to see if the module reports that it has loaded through the configuration
    # @param [String] mod module key
    # @return [Boolean] if the module key is loaded in BeEF's configuration
    def self.is_loaded(mod)
        return (self.is_enabled(mod) and BeEF::Core::Configuration.instance.get('beef.module.'+mod.to_s+'.loaded') == true)
    end

    # Returns module class definition
    # @param [String] mod module key
    # @return [Class] the module class
    def self.get_definition(mod)
        return BeEF::Core::Command.const_get(BeEF::Core::Configuration.instance.get("beef.module.#{mod.to_s}.class"))
    end

    # Gets all module options
    # @param [String] mod module key
    # @return [Hash] a hash of all the module options
    # @note API Fire: get_options 
    def self.get_options(mod)
        if BeEF::API::Registrar.instance.matched?(BeEF::API::Module, 'get_options', [mod])
            options = BeEF::API::Registrar.instance.fire(BeEF::API::Module, 'get_options', mod)
            mo = []
            options.each{|o|
                if o[:data].kind_of?(Array)
                    mo += o[:data]    
                else
                    print_debug "API Warning: return result for BeEF::Module.get_options() was not an array."
                end
            }
            return mo 
        end
        if self.check_hard_load(mod)
            class_name = BeEF::Core::Configuration.instance.get("beef.module.#{mod}.class")
            class_symbol = BeEF::Core::Command.const_get(class_name)
            if class_symbol and class_symbol.respond_to?(:options)
              return class_symbol.options
            end
            #TODO: do we really need to print this info? At then modules with no options are common,
            # so I guess we shouldn't print this info even in debug mode
#            else
#                print_debug "Module '#{mod}', no options method defined"
#            end
        end
        return []
    end

    # Gets all module payload options
    # @param [String] mod module key
    # @return [Hash] a hash of all the module options
    # @note API Fire: get_options 
    def self.get_payload_options(mod,payload)
        if BeEF::API::Registrar.instance.matched?(BeEF::API::Module, 'get_payload_options', [mod,nil])
            options = BeEF::API::Registrar.instance.fire(BeEF::API::Module, 'get_payload_options', mod,payload)
            return options 
        end
        return []
    end
    
    # Soft loads a module
    # @note A soft load consists of only loading the modules configuration (ie not the module.rb)
    # @param [String] mod module key
    # @return [Boolean] whether or not the soft load process was successful
    # @note API Fire: pre_soft_load
    # @note API Fire: post_soft_load
    def self.soft_load(mod)
        # API call for pre-soft-load module
        BeEF::API::Registrar.instance.fire(BeEF::API::Module, 'pre_soft_load', mod)
        config = BeEF::Core::Configuration.instance
        if not config.get("beef.module.#{mod}.loaded")
            if File.exists?(config.get('beef.module.'+mod+'.path')+'/module.rb')
                BeEF::Core::Configuration.instance.set('beef.module.'+mod+'.class', mod.capitalize)
                self.parse_targets(mod)
                print_debug "Soft Load module: '#{mod}'"
                # API call for post-soft-load module
                BeEF::API::Registrar.instance.fire(BeEF::API::Module, 'post_soft_load', mod)
                return true
            else
                print_debug "Unable to locate module file: #{config.get('beef.module.'+mod+'.path')}module.rb"
            end
            print_error "Unable to load module '#{mod}'"
        end
        return false
    end

    # Hard loads a module
    # @note A hard load consists of loading a pre-soft-loaded module by requiring the module.rb
    # @param [String] mod module key
    # @return [Boolean] whether or not the hard load was successful
    # @note API Fire: pre_hard_load
    # @note API Fire: post_hard_load
    def self.hard_load(mod)
        # API call for pre-hard-load module
        BeEF::API::Registrar.instance.fire(BeEF::API::Module, 'pre_hard_load', mod)
        config = BeEF::Core::Configuration.instance
        if self.is_enabled(mod)
            begin
                require config.get("beef.module.#{mod}.path")+'module.rb'
                if self.exists?(config.get("beef.module.#{mod}.class"))
                    # start server mount point
                    BeEF::Core::Server.instance.mount("/command/#{mod}.js", BeEF::Core::Handlers::Commands, mod)
                    BeEF::Core::Configuration.instance.set("beef.module.#{mod}.mount", "/command/#{mod}.js")
                    BeEF::Core::Configuration.instance.set('beef.module.'+mod+'.loaded', true)
                    print_debug "Hard Load module: '#{mod.to_s}'"
                    # API call for post-hard-load module
                    BeEF::API::Registrar.instance.fire(BeEF::API::Module, 'post_hard_load', mod)
                    return true
                else
                    print_error "Hard loaded module '#{mod.to_s}' but the class BeEF::Core::Commands::#{mod.capitalize} does not exist"
                end
            rescue => e
                BeEF::Core::Configuration.instance.set('beef.module.'+mod+'.loaded', false)
                print_error "There was a problem loading the module '#{mod.to_s}'"
                print_debug "Hard load module syntax error: #{e.to_s}"
            end
        else
            print_error "Hard load attempted on module '#{mod.to_s}' that is not enabled."
        end
        return false
    end

    # Checks to see if a module has been hard loaded, if not a hard load is attempted
    # @param [String] mod module key
    # @return [Boolean] if already hard loaded then true otherwise (see #hard_load)
    def self.check_hard_load(mod)
        if not self.is_loaded(mod)
            return self.hard_load(mod)
        end
        return true
    end
    
    # Get module key by database ID
    # @param [Integer] id module database ID
    # @return [String] module key
    def self.get_key_by_database_id(id)
        ret = BeEF::Core::Configuration.instance.get('beef.module').select {|k, v| v.has_key?('db') and v['db']['id'].to_i == id.to_i }
        return (ret.kind_of?(Array)) ? ret.first.first : ret.keys.first
    end

    # Get module key by module class
    # @param [Class] c module class
    # @return [String] module key
    def self.get_key_by_class(c)
        ret = BeEF::Core::Configuration.instance.get('beef.module').select {|k, v| v.has_key?('class') and v['class'].to_s == c.to_s }
        return (ret.kind_of?(Array)) ? ret.first.first : ret.keys.first
    end

    # Checks to see if module class exists
    # @param [String] mod module key
    # @return [Boolean] returns whether or not the class exists 
    def self.exists?(mod)
        begin
            kclass = BeEF::Core::Command.const_get(mod.capitalize)
            return kclass.is_a?(Class)
        rescue NameError
            return false
        end
    end

    # Checks target configuration to see if browser / version / operating system is supported
    # @param [String] mod module key
    # @param [Hash] opts hash of module support information
    # @return [Constant, nil] returns a resulting defined constant BeEF::Core::Constants::CommandModule::*
    # @note Support uses a rating system to provide the most accurate results.
    #   1 = All match. ie: All was defined.
    #   2 = String match. ie: Firefox was defined as working.
    #   3 = Hash match. ie: Firefox defined with 1 additional parameter (eg max_ver).
    #   4+ = As above but with extra parameters.
    #   Please note this rating system has no correlation to the return constant value BeEF::Core::Constants::CommandModule::*
    def self.support(mod, opts)
        target_config = BeEF::Core::Configuration.instance.get('beef.module.'+mod+'.target')
        if target_config and opts.kind_of? Hash
            if opts.key?('browser')
                results = []
                target_config.each{|k,m|
                    m.each{|v|
                        case v
                            when String
                                if opts['browser'] == v
                                    results << {'rating' => 2, 'const' => k}
                                end
                            when Hash
                                if opts['browser'] == v.keys.first
                                    subv = v[v.keys.first]
                                    rating = 1
                                    #version check
                                    if opts.key?('ver')
                                        if subv.key?('min_ver')
                                            if subv['min_ver'].kind_of? Fixnum and opts['ver'].to_i >= subv['min_ver']
                                                rating += 1
                                            else
                                                break
                                            end
                                        end
                                        if subv.key?('max_ver')
                                            if (subv['max_ver'].kind_of? Fixnum and opts['ver'].to_i <= subv['max_ver']) or subv['max_ver'] == "latest"
                                                rating += 1
                                            else
                                                break
                                            end
                                        end
                                    end
                                    # os check
                                    if opts.key?('os') and subv.key?('os')
                                        match = false
                                        opts['os'].each{|o|
                                            case subv['os']
                                                when String
                                                    if o == subv['os']
                                                        rating += 1
                                                        match = true
                                                    elsif subv['os'] == BeEF::Core::Constants::Os::OS_ALL_UA_STR
                                                        rating += 1
                                                        match = true
                                                    end
                                                when Array
                                                    subv['os'].each{|p|
                                                        if o == p or p == BeEF::Core::Constants::Os::OS_ALL_UA_STR
                                                            rating += 1
                                                            match = true
                                                        end
                                                    }
                                            end
                                        }
                                        if not match 
                                            break
                                        end
                                    end
                                    if rating != 1
                                        results << {'rating' => rating, 'const' => k}
                                    end
                                end
                        end
                        if v == BeEF::Core::Constants::Browsers::ALL
                            results << {'rating' => 1, 'const' => k}
                        end
                    }
                }
                if results.count > 0
                    return results.sort_by {|v| v['rating']}.last['const']
                else
                    return BeEF::Core::Constants::CommandModule::VERIFIED_UNKNOWN
                end
            else
                print_error "BeEF::Module.support() was passed a hash without a valid browser constant"
            end
        end
        return nil
    end

    # Translates module target configuration
    # @note Takes the user defined target configuration and replaces it with equivalent a constant based generated version
    # @param [String] mod module key
    def self.parse_targets(mod)
        target_config = BeEF::Core::Configuration.instance.get('beef.module.'+mod+'.target')
        if target_config
            targets = {}
            target_config.each{|k,v|
                begin
                    if BeEF::Core::Constants::CommandModule.const_defined?('VERIFIED_'+k.upcase)
                        key = BeEF::Core::Constants::CommandModule.const_get('VERIFIED_'+k.upcase)
                        if not targets.key?(key)
                            targets[key] = []
                        end
                        browser = nil
                        case v
                            when String
                                browser = self.match_target_browser(v)
                                if browser
                                    targets[key] << browser
                                end
                            when Array
                                v.each{|c|
                                    browser = self.match_target_browser(c)
                                    if browser
                                        targets[key] << browser
                                    end
                                }
                            when Hash
                                v.each{|k,c|
                                    browser = self.match_target_browser(k)
                                    if browser
                                        case c
                                            when TrueClass
                                                targets[key] << browser
                                            when Hash
                                                details = self.match_target_browser_spec(c)
                                                if details
                                                    targets[key] << {browser => details}
                                                end
                                        end
                                    end
                                }
                        end
                    end
                rescue NameError
                    print_error "Module \"#{mod}\" configuration has invalid target status defined \"#{k}\""
                end
            }
            BeEF::Core::Configuration.instance.clear("beef.module.#{mod}.target")
            BeEF::Core::Configuration.instance.set("beef.module.#{mod}.target", targets)
        end
    end

    # Translates simple browser target configuration
    # @note Takes a user defined browser type and translates it into a BeEF constant
    # @param [String] v user defined browser
    # @return [Constant] a BeEF browser constant
    def self.match_target_browser(v)
        browser = false
        if v.class == String
            begin
                if BeEF::Core::Constants::Browsers.const_defined?(v.upcase)
                    browser = BeEF::Core::Constants::Browsers.const_get(v.upcase)
                end
            rescue NameError
                print_error "Could not identify browser target specified as \"#{v}\""
            end
        else
            print_error "Invalid datatype passed to BeEF::Module.match_target_browser()"
        end
        return browser
    end

    # Translates complex browser target configuration
    # @note Takes a complex user defined browser hash and converts it to applicable BeEF constants
    # @param [Hash] v user defined browser hash    
    # @return [Hash] BeEF constants hash
    def self.match_target_browser_spec(v)
        browser = {}
        if v.class == Hash
            if v.key?('max_ver') and (v['max_ver'].is_a?(Fixnum) or v['max_ver'].is_a?(Float) or v['max_ver'] == "latest")
                browser['max_ver'] = v['max_ver']
            end
            if v.key?('min_ver') and (v['min_ver'].is_a?(Fixnum) or v['min_ver'].is_a?(Float))
                browser['min_ver'] = v['min_ver']
            end
            if v.key?('os')
                case v['os']
                    when String
                        os = self.match_target_os(v['os'])
                        if os 
                            browser['os'] = os 
                        end
                    when Array
                        browser['os'] = []
                        v['os'].each{|c|
                            os = self.match_target_os(c)
                            if os
                                browser['os'] << os 
                            end
                        }
                end
            end
        else
            print_error "Invalid datatype passed to BeEF::Module.match_target_browser_spec()"
        end
        return browser
    end

    # Translates simple OS target configuration
    # @note Takes user defined OS specification and translates it into BeEF constants
    # @param [String] v user defined OS string
    # @return [Constant] BeEF OS Constant
    def self.match_target_os(v)
        os = false
        if v.class == String
            begin
                if BeEF::Core::Constants::Os.const_defined?("OS_#{v.upcase}_UA_STR")
                    os = BeEF::Core::Constants::Os.const_get("OS_#{v.upcase}_UA_STR")
                end
            rescue NameError
                print_error "Could not identify OS target specified as \"#{v}\""
            end
        else
            print_error "Invalid datatype passed to BeEF::Module.match_target_os()"
        end
        return os
    end

    # Executes a module 
    # @param [String] mod module key
    # @param [String] hbsession hooked browser session
    # @param [Array] opts array of module execute options (see #get_options)
    # @return [Boolean] whether or not the BeEF system executed the module
    # @note The return value of this function does not specify if the module was successful, only that it was executed within the framework
    def self.execute(mod, hbsession, opts=[])
        if not (self.is_present(mod) and self.is_enabled(mod))
            print_error "Module not found '#{mod}'. Failed to execute module."
            return false
        end
        if BeEF::API::Registrar.instance.matched?(BeEF::API::Module, 'override_execute', [mod, nil,nil])
            BeEF::API::Registrar.instance.fire(BeEF::API::Module, 'override_execute', mod, hbsession,opts)
            # @note We return true by default as we cannot determine the correct status if multiple API hooks have been called
            return true
        end
        hb = BeEF::HBManager.get_by_session(hbsession)        
        if not hb
            print_error "Could not find hooked browser when attempting to execute module '#{mod}'"
            return false
        end
        self.check_hard_load(mod)
        command_module = self.get_definition(mod).new(mod)
        if command_module.respond_to?(:pre_execute)
            command_module.pre_execute
        end
        h = self.merge_options(mod, [])
        c = BeEF::Core::Models::Command.new(:data => self.merge_options(mod, opts).to_json,
            :hooked_browser_id => hb.id,
            :command_module_id => BeEF::Core::Configuration.instance.get("beef.module.#{mod}.db.id"),
            :creationdate => Time.new.to_i
          ).save
        return true
    end

    # Merges default module options with array of custom options
    # @param [String] mod module key
    # @param [Hash] h module options customised by user input
    # @return [Hash, nil] returns merged options
    def self.merge_options(mod, h)
        if self.is_present(mod)
            self.check_hard_load(mod)
            merged = []
            defaults = self.get_options(mod)
            defaults.each{|v|
                mer = nil
                h.each{|o|
                    if v.has_key?('name') and o.has_key?('name') and v['name'] == o['name']
                        mer = v.deep_merge(o)
                    end
                }
                if mer != nil
                    merged.push(mer)
                else
                    merged.push(v)
                end
            }
            return merged
        end
        return nil
    end

end
end


