# frozen_string_literal: true

require 'ipaddr'

module Msf::MCP
  module Security
    class InputValidator
      LIMIT_DEFAULT = 100
      LIMIT_MIN = 1
      LIMIT_MAX = 1000

      # Generic parameter validation against a constraint
      #
      # Dispatches based on the constraint type:
      # - Array  → value must be included in the list (enum)
      # - Range  → value must be an integer within the range, or a Range whose
      #            bounds are within the constraint (range must be integer-bounded)
      # - Regexp → value (via .to_s) must match the pattern
      #
      # @param name [String] Parameter name (used in error messages)
      # @param value [Object] Value to validate
      # @param constraint [Array, Range, Regexp] Allowed values, range, or pattern
      # @param allow_nil [Boolean] Whether nil/empty values are allowed (default: false)
      # @param max_size [Integer] (optional) Maximum length for string values (only applies to Regexp constraints)
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_parameter!(name, value, constraint, allow_nil: false, max_size: nil)
        if allow_nil
          return true if value.nil?
          return true if value.respond_to?(:empty?) && value.empty?
        else
          raise ValidationError, "#{name} cannot be nil" if value.nil?
          raise ValidationError, "#{name} cannot be empty" if value.respond_to?(:empty?) && value.empty?
        end

        case constraint
        when Array
          unless constraint.include?(value)
            raise ValidationError, "Invalid #{name}: #{value.inspect}. Must be one of: #{constraint.join(', ')}"
          end
        when Range
          unless constraint.first.is_a?(Integer) && constraint.last.is_a?(Integer)
            raise ArgumentError, "Range constraint must be a range of integers, got #{constraint.first.class}..#{constraint.last.class}"
          end
          if value.is_a?(Range)
            begin
              int_first = Integer(value.first)
              int_last = Integer(value.last)
            rescue TypeError, ArgumentError
              raise ValidationError, "#{name} must have integer bounds: #{value.inspect}"
            end
            unless constraint.cover?(int_first..int_last)
              raise ValidationError, "#{name} must be between #{constraint.min} and #{constraint.max}: #{int_first}..#{int_last}"
            end
          else
            begin
              int_value = Integer(value)
            rescue TypeError, ArgumentError
              raise ValidationError, "#{name} must be an integer: #{value.inspect}"
            end
            unless constraint.cover?(int_value)
              raise ValidationError, "#{name} must be between #{constraint.min} and #{constraint.max}: #{value}"
            end
          end
        when Regexp
          string_value = value.to_s
          if max_size && string_value.length > max_size
            raise ValidationError, "#{name} too long (max #{max_size} characters)"
          end
          unless string_value.match?(constraint)
            raise ValidationError, "Invalid #{name} format: #{value}"
          end
        else
          raise ArgumentError, "Unsupported constraint type: #{constraint.class}"
        end

        true
      end

      # Validate IP address or CIDR range
      #
      # @param addr [String] IP address or CIDR (e.g., "192.168.1.1" or "192.168.1.0/24")
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_ip_address!(addr)
        return true if addr.nil? || addr.empty?

        begin
          IPAddr.new(addr)
          true
        rescue IPAddr::InvalidAddressError
          raise ValidationError, "Invalid IP address or CIDR: #{addr}"
        end
      end

      # Validate port or port range
      #
      # @param range [String, Integer] Port number or range (e.g., "80" or "80-443")
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_port_range!(range)
        return true if range.nil? || range.to_s.empty?

        range_str = range.to_s

        # Match a port range like "80-443" — requires digits on both sides of the dash
        if range_str.match?(/\A\s*[[:alnum:]]+-[[:alnum:]]+\s*\z/)
          begin
            start_port, end_port = range_str.split('-', 2).map { |p| Integer(p.strip) }
          rescue TypeError, ArgumentError
            raise ValidationError, "Port range must have integer bounds: #{range_str}"
          end
          validate_parameter!('Port range', start_port..end_port, 1..65535)
        else
          validate_parameter!('Port', range_str, 1..65535)
        end

        true
      end

      # Validate query string for module search
      #
      # @param query [String] Search query
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_search_query!(query)
        validate_parameter!('Search query', query, /\A[[:print:]]+\z/, allow_nil: false, max_size: 500)
      end

      # Validate limit parameter for pagination
      #
      # @param limit [Integer] Limit value
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_limit!(limit)
        validate_parameter!('Limit', limit, LIMIT_MIN..LIMIT_MAX, allow_nil: true)
      end

      # Validate offset parameter for pagination
      #
      # @param offset [Integer] Offset value
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_offset!(offset)
        validate_parameter!('Offset', offset, 0..LIMIT_MAX, allow_nil: true)
      end

      # Validate pagination parameters
      #
      # @param limit [Integer] Limit value
      # @param offset [Integer] Offset value
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_pagination!(limit, offset)
        validate_limit!(limit)
        validate_offset!(offset)
      end

      # Validate module type
      #
      # @param module_type [String] Module type
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_module_type!(module_type)
        validate_parameter!('Module type', module_type, %w[exploit auxiliary post payload encoder evasion nop])
      end

      # Validate module name
      #
      # @param module_name [String] Module name/path
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_module_name!(module_name)
        # Basic path validation (alphanumeric, slashes, underscores, hyphens)
        validate_parameter!('Module name', module_name, %r{\A[\w/\-]+\z}, max_size: 500)
      end

      # Validate only_up boolean parameter
      #
      # @param only_up [Boolean] Only up parameter
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_only_up!(only_up)
        validate_parameter!('only_up', only_up, [true, false])
      end

      # Validate protocol parameter
      #
      # @param protocol [String] Protocol ('tcp' or 'udp')
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_protocol!(protocol)
        validate_parameter!('Protocol', protocol.to_s.downcase, %w[tcp udp], allow_nil: true)
      end

      # Module datastore option limits.
      #
      # Bounds an oversized options hash before it reaches the Metasploit
      # datastore. Real modules register on the order of 10-60 options across
      # all inherited mixins, so 100 is a comfortable ceiling. This is a
      # defence-in-depth cap against DoS-style input abuse from an MCP client,
      # not a transport limit (MCP itself does not restrict payload size).
      MODULE_OPTIONS_MAX_KEYS = 100
      # Matches the option name conventions used across the Framework: standard
      # SCREAMING_SNAKE_CASE (RHOSTS, LPORT), CamelCase advanced options
      # (SSLVersion, BasicAuth), namespaced mixin options separated by `::`
      # (HTTP::compression, CMDSTAGER::FLAVOR, EXE::Custom), multi-level
      # namespaces (HTML::javascript::escape), and hyphenated header-style
      # names (BEARER-TOKEN). Each segment must start with a letter or
      # underscore. Leading/trailing separators and double separators
      # (`:::`, `--`) are rejected.
      MODULE_OPTIONS_KEY_REGEX = /\A[A-Za-z_][A-Za-z0-9_]*(?:(?:::|-)[A-Za-z_][A-Za-z0-9_]*)*\z/
      MODULE_OPTIONS_KEY_MAX_LENGTH = 128
      MODULE_OPTIONS_VALUE_MAX_BYTES = 8 * 1024
      MODULE_OPTIONS_TOTAL_MAX_BYTES = 64 * 1024

      # Allowed scalar types for datastore values. The Metasploit datastore is flat,
      # so no nested Hashes / Arrays are permitted.
      MODULE_OPTION_SCALAR_TYPES = [String, Integer, Float, TrueClass, FalseClass, NilClass].freeze

      # UUID format produced by Metasploit's RPC layer via Rex::Text.rand_text_alphanumeric(24).
      # The generator emits exactly 24 characters drawn from a mixed-case alphanumeric pool.
      MODULE_RUN_UUID_REGEX = /\A[A-Za-z0-9]{24}\z/

      SESSION_ID_RANGE = 1..65535
      # Character cap on a single session.interactive_write payload.
      SESSION_DATA_MAX_CHARS = 10_000

      # Validate module datastore options hash.
      #
      # Enforces structural limits before any RPC call so a misbehaving client
      # cannot push an unbounded payload through to the framework.
      #
      # Keys may be Strings or Symbols (the MCP transport deep-symbolizes JSON
      # input). Either form is validated against MODULE_OPTIONS_KEY_REGEX via
      # `to_s`; the tool layer is responsible for normalizing to String keys
      # before forwarding to the Metasploit datastore.
      #
      # @param options [Hash] Module datastore options
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_module_options!(options)
        raise ValidationError, 'Module options must be a Hash' unless options.is_a?(Hash)

        if options.size > MODULE_OPTIONS_MAX_KEYS
          raise ValidationError, "Module options has too many keys (max #{MODULE_OPTIONS_MAX_KEYS})"
        end

        total_bytes = 0
        options.each do |key, value|
          unless key.is_a?(String) || key.is_a?(Symbol)
            raise ValidationError, "Invalid module option key: #{key.inspect}"
          end

          key_str = key.to_s
          if key_str.length > MODULE_OPTIONS_KEY_MAX_LENGTH
            raise ValidationError, "Module option key exceeds #{MODULE_OPTIONS_KEY_MAX_LENGTH} chars: #{key.inspect}"
          end
          unless key_str.match?(MODULE_OPTIONS_KEY_REGEX)
            raise ValidationError, "Invalid module option key: #{key.inspect}"
          end

          unless MODULE_OPTION_SCALAR_TYPES.any? { |klass| value.is_a?(klass) }
            raise ValidationError, "Invalid module option value for #{key}: must be a scalar (String, Integer, Float, Boolean, or nil)"
          end

          if value.is_a?(String) && value.bytesize > MODULE_OPTIONS_VALUE_MAX_BYTES
            raise ValidationError, "Module option #{key} string value exceeds #{MODULE_OPTIONS_VALUE_MAX_BYTES} bytes"
          end

          total_bytes += key.to_s.bytesize
          total_bytes += value.bytesize if value.is_a?(String)
          if total_bytes > MODULE_OPTIONS_TOTAL_MAX_BYTES
            raise ValidationError, "Module options total payload exceeds #{MODULE_OPTIONS_TOTAL_MAX_BYTES} bytes"
          end
        end

        true
      end

      # Validate a module run UUID returned by module.execute / module.check.
      #
      # @param uuid [String] UUID string
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_uuid!(uuid)
        raise ValidationError, 'UUID must be a String' unless uuid.is_a?(String)

        validate_parameter!('UUID', uuid, MODULE_RUN_UUID_REGEX, max_size: 24)
      end

      # Validate a session identifier.
      #
      # @param session_id [Integer] Session ID
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_session_id!(session_id)
        raise ValidationError, 'Session ID must be an Integer' unless session_id.is_a?(Integer)

        validate_parameter!('Session ID', session_id, SESSION_ID_RANGE)
      end

      # Validate session input data for session.interactive_write.
      #
      # @param data [String] Data to send to the session
      # @return [true] If valid
      # @raise [ValidationError] If invalid
      def self.validate_session_data!(data)
        raise ValidationError, 'Session data must be a String' unless data.is_a?(String)
        raise ValidationError, 'Session data cannot be empty' if data.empty?

        if data.length > SESSION_DATA_MAX_CHARS
          raise ValidationError, "Session data exceeds #{SESSION_DATA_MAX_CHARS} characters"
        end

        true
      end
    end
  end
end
