require 'rex/socket'

module Msf::Exploit::Remote::Relay
  # A thread safe target list. The provided targets will be iterated over via the {next} method.
  class TargetList
    include MonitorMixin

    # @param [String] targets
    def initialize(protocol, port, targets, path=nil, randomize_targets: true, drop_mic_only: false, drop_mic_and_sign_key_exch_flags: false, protocol_options: {})
      super()

      targets = Rex::Socket::RangeWalker.new(targets).to_enum(:each_ip).map do |target_ip|
        Target.new(
          ip: target_ip,
          port: port,
          protocol: protocol,
          path: path,
          drop_mic_only: drop_mic_only,
          drop_mic_and_sign_key_exch_flags: drop_mic_and_sign_key_exch_flags,
          protocol_options: protocol_options
        )
      end
      @targets = randomize_targets ? targets.shuffle : targets
    end

    def each(&block)
      @targets.each(&block)
    end

    # Return the next available target, or nil if the identity has been relayed against all targets
    # @param [String,nil] identity The identity, i.e. domain/user, if available
    # @return [Target,nil] The next target for the given identity with the least amount of relay attempts. Or nil if all targets have been relayed to for that identity
    def next(identity)
      synchronize do
        next_target = next_target_for(identity)
        return nil if next_target.nil?

        next_target.on_relay_start(identity)
        next_target
      end
    end

    # Updates tracking to mark a host as being successfully relayed or not
    # @param [Msf::Exploit::Remote::SMB::Relay::Target] target The target that was successfully relayed or not
    # @param [String] identity the identity which was used as part of relaying
    # @param [TrueClass|FalseClass] is_success True when this identity was successfully relayed to the target, false otherwise
    def on_relay_end(target, identity:, is_success:)
      synchronize do
        target.on_relay_end(identity: identity, is_success: is_success)
      end
    end

    private

    # @param [Object] identity The identity that will be used during the relay process
    # @return [Enumerator<Object>] All targets that have not yet successfully been relayed to a target based on the given identity
    def next_target_for(identity)
      # Choose the next target that hasn't been relayed to yet, and has the least retry attempts - in order to try and
      # round robin requests to each host
      next_target = @targets.select { |target| target.eligible_relay_target?(identity) }
                            .min_by { |target| target.relay_attempts_for(identity) }

      next_target
    end
  end

  class Target
    def initialize(ip:, port:, protocol:, path: nil, drop_mic_only: false, drop_mic_and_sign_key_exch_flags: false, protocol_options: nil)
      @ip = ip
      @port = port
      @protocol = protocol
      @path = path
      @protocol_options = protocol_options
      @relay_state = Hash.new do |hash, identity|
        hash[identity] = {
          relay_status: nil,
          relay_attempted_at: nil,
          relayed_at: nil,
          relay_attempts: 0
        }
      end
      @drop_mic_only= drop_mic_only
      @drop_mic_and_sign_key_exch_flags = drop_mic_and_sign_key_exch_flags
    end

    attr_reader :ip, :port, :protocol, :path, :drop_mic_only, :drop_mic_and_sign_key_exch_flags, :protocol_options

    def eligible_relay_target?(identity)
      return true if identity.nil?

      relay_data = relay_data_for(identity)
      relay_data[:relay_status].nil? || relay_data[:relay_status] == :failed
    end

    def relay_attempts_for(identity)
      relay_data = relay_data_for(identity)
      relay_data[:relay_attempts]
    end

    def on_relay_start(identity)
      relay_data = relay_data_for(identity)
      relay_data[:relay_attempts] += 1
      relay_data[:relay_attempted_at] = Time.now
      relay_data[:relay_status] = :attempting
    end

    def on_relay_end(identity:, is_success:)
      relay_data = relay_data_for(identity)
      if is_success
        relay_data[:relay_status] = :success
        relay_data[:relayed_at] = Time.now
      else
        relay_data[:relay_status] = :failed
      end
    end

    def to_h
      { ip: ip, port: port, protocol: protocol, path: path, relay_state: @relay_state }
    end

    def to_s
      s = "#{protocol}://#{Rex::Socket.to_authority(ip, port)}"
      s << ('/' + path.delete_prefix('/')) unless path.blank?
      s
    end

    private

    def relay_data_for(username)
      @relay_state[username]
    end
  end
end
