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

3
module Rex
4
module Ui
5
module Text
6

7
begin
8

9
  ###
10
  #
11
  # This class implements standard input using readline against
12
  # standard input.  It supports tab completion.
13
  #
14
  ###
15
  class Input::Readline < Rex::Ui::Text::Input
16

17
    #
18
    # Initializes the readline-aware Input instance for text.
19
    #
20
    def initialize(tab_complete_proc = nil)
21
      if(not Object.const_defined?('Readline'))
22
        require 'readline'
23
      end
24

25
      self.extend(::Readline)
26

27
      if tab_complete_proc
28
        ::Readline.basic_word_break_characters = ""
29
        @rl_saved_proc = with_error_handling(tab_complete_proc)
30
        ::Readline.completion_proc = @rl_saved_proc
31
      end
32
    end
33

34
    #
35
    # Reattach the original completion proc
36
    #
37
    def reset_tab_completion(tab_complete_proc = nil)
38
      ::Readline.basic_word_break_characters = "\x00"
39
      ::Readline.completion_proc = tab_complete_proc ? with_error_handling(tab_complete_proc) : @rl_saved_proc
40
    end
41

42

43
    #
44
    # Retrieve the line buffer
45
    #
46
    def line_buffer
47
      if defined? RbReadline
48
        RbReadline.rl_line_buffer
49
      else
50
        ::Readline.line_buffer
51
      end
52
    end
53

54
    attr_accessor :prompt
55

56
    #
57
    # Whether or not the input medium supports readline.
58
    #
59
    def supports_readline
60
      true
61
    end
62

63
    #
64
    # Calls sysread on the standard input handle.
65
    #
66
    def sysread(len = 1)
67
      begin
68
        self.fd.sysread(len)
69
      rescue ::Errno::EINTR
70
        retry
71
      end
72
    end
73

74
    #
75
    # Read a line from stdin
76
    #
77
    def gets()
78
      begin
79
        self.fd.gets()
80
      rescue ::Errno::EINTR
81
        retry
82
      end
83
    end
84

85
    #
86
    # Stick readline into a low-priority thread so that the scheduler doesn't slow
87
    # down other background threads. This is important when there are many active
88
    # background jobs, such as when the user is running Karmetasploit
89
    #
90
    def pgets
91

92
      line = nil
93
      orig = Thread.current.priority
94

95
      begin
96
        Thread.current.priority = -20
97

98
        output.prompting
99
        line = readline_with_output(prompt, true)
100
        ::Readline::HISTORY.pop if (line and line.empty?)
101
      ensure
102
        Thread.current.priority = orig || 0
103
        output.prompting(false)
104
      end
105

106
      line
107
    end
108

109
    #
110
    # Returns the output pipe handle
111
    #
112
    def fd
113
      $stdin
114
    end
115

116
    #
117
    # Indicates that this input medium as a shell builtin, no need
118
    # to extend.
119
    #
120
    def intrinsic_shell?
121
      true
122
    end
123

124
    #
125
    # The prompt that is to be displayed.
126
    #
127
    attr_accessor :prompt
128
    #
129
    # The output handle to use when displaying the prompt.
130
    #
131
    attr_accessor :output
132

133
    private
134

135
    def readline_with_output(prompt, add_history=false)
136
      # rb-readlines's Readline.readline hardcodes the input and output to
137
      # $stdin and $stdout, which means setting `Readline.input` or
138
      # `Readline.ouput` has no effect when running `Readline.readline` with
139
      # rb-readline, so need to reimplement
140
      # []`Readline.readline`](https://github.com/luislavena/rb-readline/blob/ce4908dae45dbcae90a6e42e3710b8c3a1f2cd64/lib/readline.rb#L36-L58)
141
      # for rb-readline to support setting input and output.  Output needs to
142
      # be set so that colorization works for the prompt on Windows.
143
      self.prompt = prompt
144

145
      # TODO: there are unhandled quirks in async output buffering that
146
      # we have not solved yet, for instance when loading meterpreter
147
      # extensions, supporting Windows, printing output from commands, etc.
148
      # Remove this guard when issues are resolved.
149
=begin
150
      reset_sequence = "\n\001\r\033[K\002"
151
      if (/mingw/ =~ RUBY_PLATFORM)
152
        reset_sequence = ""
153
      end
154
=end
155
      reset_sequence = ""
156

157
      if defined? RbReadline
158
        RbReadline.rl_instream = fd
159
        RbReadline.rl_outstream = output
160

161
        begin
162
          line = RbReadline.readline(reset_sequence + prompt)
163
        rescue ::Exception => exception
164
          RbReadline.rl_cleanup_after_signal()
165
          RbReadline.rl_deprep_terminal()
166

167
          raise exception
168
        end
169

170
        if add_history && line && !line.start_with?(' ')
171
          # Don't add duplicate lines to history
172
          if ::Readline::HISTORY.empty? || line.strip != ::Readline::HISTORY[-1]
173
            RbReadline.add_history(line.strip)
174
          end
175
        end
176

177
        line.try(:dup)
178
      else
179
        # The line that's read is immediately added to history
180
        line = ::Readline.readline(reset_sequence + prompt, true)
181

182
        # Don't add duplicate lines to history
183
        if ::Readline::HISTORY.length > 1 && line == ::Readline::HISTORY[-2]
184
          ::Readline::HISTORY.pop
185
        end
186

187
        line
188
      end
189
    end
190

191
    private
192

193
    def with_error_handling(proc)
194
      proc do |*args|
195
        proc.call(*args)
196
      rescue StandardError => e
197
        elog("tab_complete_proc has failed with args #{args}", error: e)
198
        []
199
      end
200
    end
201

202
  end
203
rescue LoadError
204
end
205

206
end
207
end
208
end
209

210