#
# Copyright (c) 2006-2024 Wade Alcorn - wade@bindshell.net
# Browser Exploitation Framework (BeEF) - https://beefproject.com
# See the file 'doc/COPYING' for copying permission
#
module BeEF
  module Core
    module Models
      module Dns
        # Represents an individual DNS rule.
        class Rule < BeEF::Core::Model
          # Hooks the model's "save" event. Validates pattern/response and generates a rule identifier.
          before_save :check_rule
          self.table_name = 'dns_rules'
          serialize :response, Array

          private

          def check_rule
            validate_pattern(pattern)
            self.callback = format_callback(resource.constantize, response)
          rescue InvalidDnsPatternError, UnknownDnsResourceError, InvalidDnsResponseError => e
            print_error e.message
            throw :halt

            # self.id = BeEF::Core::Crypto.dns_rule_id
          end

          # Verifies that the given pattern is valid (i.e. non-empty, no null's or printable characters).
          def validate_pattern(pattern)
            raise InvalidDnsPatternError unless BeEF::Filters.is_non_empty_string?(pattern) &&
                                                !BeEF::Filters.has_null?(pattern) &&
                                                !BeEF::Filters.has_non_printable_char?(pattern)
          end

          # Strict validator which ensures that only an appropriate response is given.
          #
          # @param resource [Resolv::DNS::Resource::IN] resource record type
          # @param response [String, Symbol, Array] response to include in callback
          #
          # @return [String] string representation of callback that can safely be eval'd
          def format_callback(resource, response)
            sym_regex = /^:?(NoError|FormErr|ServFail|NXDomain|NotImp|Refused|NotAuth)$/i

            if resource == Resolv::DNS::Resource::IN::A
              if response.is_a?(String) && BeEF::Filters.is_valid_ip?(response, :ipv4)
                format "t.respond!('%s')", response
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response.to_s =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              elsif response.is_a?(Array)
                str1 = "t.respond!('%s');"
                str2 = ''

                response.each do |r|
                  raise InvalidDnsResponseError, 'A' unless BeEF::Filters.is_valid_ip?(r, :ipv4)

                  str2 << format(str1, r)
                end

                str2
              else
                raise InvalidDnsResponseError, 'A'
              end
            elsif resource == Resolv::DNS::Resource::IN::AAAA
              if response.is_a?(String) && BeEF::Filters.is_valid_ip?(response, :ipv6)
                format "t.respond!('%s')", response
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              elsif response.is_a?(Array)
                str1 = "t.respond!('%s');"
                str2 = ''

                response.each do |r|
                  raise InvalidDnsResponseError, 'AAAA' unless BeEF::Filters.is_valid_ip?(r, :ipv6)

                  str2 << format(str1, r)
                end

                str2
              else
                raise InvalidDnsResponseError, 'AAAA'
              end
            elsif resource == Resolv::DNS::Resource::IN::CNAME
              if response.is_a?(String) && BeEF::Filters.is_valid_domain?(response)
                format "t.respond!(Resolv::DNS::Name.create('%s'))", response
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              else
                raise InvalidDnsResponseError, 'CNAME'
              end
            elsif resource == Resolv::DNS::Resource::IN::MX
              if response[0].is_a?(Integer) &&
                 BeEF::Filters.is_valid_domain?(response[1])

                data = { preference: response[0], exchange: response[1] }
                format "t.respond!(%<preference>d, Resolv::DNS::Name.create('%<exchange>s'))", data
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              else
                raise InvalidDnsResponseError, 'MX'
              end
            elsif resource == Resolv::DNS::Resource::IN::NS
              if response.is_a?(String) && BeEF::Filters.is_valid_domain?(response)
                format "t.respond!(Resolv::DNS::Name.create('%s'))", response
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              elsif response.is_a?(Array)
                str1 = "t.respond!(Resolv::DNS::Name.create('%s'))"
                str2 = ''

                response.each do |r|
                  raise InvalidDnsResponseError, 'NS' unless BeEF::Filters.is_valid_domain?(r)

                  str2 << format(str1, r)
                end

                str2
              else
                raise InvalidDnsResponseError, 'NS'
              end
            elsif resource == Resolv::DNS::Resource::IN::PTR
              if response.is_a?(String) && BeEF::Filters.is_valid_domain?(response)
                format "t.respond!(Resolv::DNS::Name.create('%s'))", response
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              else
                raise InvalidDnsResponseError, 'PTR'
              end
            elsif resource == Resolv::DNS::Resource::IN::SOA
              if response.is_a?(Array)
                unless BeEF::Filters.is_valid_domain?(response[0]) &&
                       BeEF::Filters.is_valid_domain?(response[1]) &&
                       response[2].is_a?(Integer) &&
                       response[3].is_a?(Integer) &&
                       response[4].is_a?(Integer) &&
                       response[5].is_a?(Integer) &&
                       response[6].is_a?(Integer)

                  raise InvalidDnsResponseError, 'SOA'
                end

                data = {
                  mname: response[0],
                  rname: response[1],
                  serial: response[2],
                  refresh: response[3],
                  retry: response[4],
                  expire: response[5],
                  minimum: response[6]
                }

                format "t.respond!(Resolv::DNS::Name.create('%<mname>s'), " +
                       "Resolv::DNS::Name.create('%<rname>s'), " +
                       '%<serial>d, ' +
                       '%<refresh>d, ' +
                       '%<retry>d, ' +
                       '%<expire>d, ' +
                       '%<minimum>d)',
                       data
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              else
                raise InvalidDnsResponseError, 'SOA'
              end
            elsif resource == Resolv::DNS::Resource::IN::WKS
              if response.is_a?(Array)
                if !BeEF::Filters.is_valid_ip?(resource[0]) &&
                   resource[1].is_a?(Integer) &&
                   resource[2].is_a?(Integer) && !resource.is_a?(String)
                  raise InvalidDnsResponseError, 'WKS'
                end

                data = {
                  address: response[0],
                  protocol: response[1],
                  bitmap: response[2]
                }

                format "t.respond!('%<address>s', %<protocol>d, %<bitmap>d)", data
              elsif (response.is_a?(Symbol) && response.to_s =~ sym_regex) || response =~ sym_regex
                format 't.fail!(:%s)', response.to_sym
              else
                raise InvalidDnsResponseError, 'WKS'
              end
            else
              raise UnknownDnsResourceError
            end
          end

          # Raised when an invalid pattern is given.
          class InvalidDnsPatternError < StandardError
            DEFAULT_MESSAGE = 'Failed to add DNS rule with invalid pattern'

            def initialize(message = nil)
              super(message || DEFAULT_MESSAGE)
            end
          end

          # Raised when a response is not valid for the given DNS resource record.
          class InvalidDnsResponseError < StandardError
            def initialize(message = nil)
              str = 'Failed to add DNS rule with invalid response for %s resource record', message
              message = format str, message unless message.nil?
              super(message)
            end
          end

          # Raised when an unknown DNS resource record is given.
          class UnknownDnsResourceError < StandardError
            DEFAULT_MESSAGE = 'Failed to add DNS rule with unknown resource record'

            def initialize(message = nil)
              super(message || DEFAULT_MESSAGE)
            end
          end
        end
      end
    end
  end
end