CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
# -*- coding: binary -*-
2

3
require 'timeout'
4
require 'thread'
5

6
module Rex
7
module Post
8
module Meterpreter
9

10
###
11
#
12
# This class handles waiting for a response to a given request
13
# and the subsequent response association.
14
#
15
###
16
class PacketResponseWaiter
17

18
  # Arbitrary argument to {#completion_routine}
19
  #
20
  # @return [Object,nil]
21
  attr_accessor :completion_param
22

23
  # A callback to be called when this waiter is notified of a packet's
24
  # arrival. If not nil, this will be called with the response packet as first
25
  # parameter and {#completion_param} as the second.
26
  #
27
  # @return [Proc,nil]
28
  attr_accessor :completion_routine
29

30
  # @return [ConditionVariable]
31
  attr_accessor :cond
32

33
  # @return [Mutex]
34
  attr_accessor :mutex
35

36
  # @return [Packet]
37
  attr_accessor :response
38

39
  # @return [Integer] request ID to wait for
40
  attr_accessor :rid
41

42
  #
43
  # Initializes a response waiter instance for the supplied request
44
  # identifier.
45
  #
46
  def initialize(rid, completion_routine = nil, completion_param = nil)
47
    self.rid      = rid.dup
48
    self.response = nil
49

50
    if (completion_routine)
51
      self.completion_routine = completion_routine
52
      self.completion_param   = completion_param
53
    else
54
      self.mutex = Mutex.new
55
      self.cond  = ConditionVariable.new
56
    end
57
  end
58

59
  #
60
  # Checks to see if this waiter instance is waiting for the supplied
61
  # packet based on its request identifier.
62
  #
63
  def waiting_for?(packet)
64
    return (packet.rid == rid)
65
  end
66

67
  #
68
  # Notifies the waiter that the supplied response packet has arrived.
69
  #
70
  # @param response [Packet]
71
  # @return [void]
72
  def notify(response)
73
    if (self.completion_routine)
74
      self.response = response
75
      self.completion_routine.call(response, self.completion_param)
76
    else
77
      self.mutex.synchronize do
78
        self.response = response
79
        self.cond.signal
80
      end
81
    end
82
  end
83

84
  #
85
  # Wait for a given time interval for the response packet to arrive.
86
  #
87
  # @param interval [Integer,nil] number of seconds to wait, or nil to wait
88
  #   forever
89
  # @return [Packet,nil] the response, or nil if the interval elapsed before
90
  #   receiving one
91
  def wait(interval)
92
    interval = nil if interval and interval == -1
93
    self.mutex.synchronize do
94
      if self.response.nil?
95
        self.cond.wait(self.mutex, interval)
96
      end
97
    end
98
    return self.response
99
  end
100

101
end
102

103
end; end; end
104

105

106