Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

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

3
require 'singleton'
4

5
module Rex
6
module Ui
7
module Text
8
module Shell
9

10
class HistoryManager
11

12
  MAX_HISTORY = 2000
13

14
  def initialize
15
    @contexts = []
16
    @debug = false
17
    # Values dequeued before work is started
18
    @write_queue = ::Queue.new
19
    # Values dequeued after work is completed
20
    @remaining_work = ::Queue.new
21
  end
22

23
  # Create a new history command context when executing the given block
24
  #
25
  # @param [String,nil] history_file The file to load and persist commands to
26
  # @param [String] name Human readable history context name
27
  # @param [Symbol] input_library The input library to provide context for. :reline, :readline
28
  # @param [Proc] block
29
  # @return [nil]
30
  def with_context(history_file: nil, name: nil, input_library: nil, &block)
31
    # Default to Readline for backwards compatibility.
32
    push_context(history_file: history_file, name: name, input_library: input_library || :readline)
33

34
    begin
35
      block.call
36
    ensure
37
      pop_context
38
    end
39

40
    nil
41
  end
42

43
  # Flush the contents of the write queue to disk. Blocks synchronously.
44
  def flush
45
    until @write_queue.empty? && @remaining_work.empty?
46
      sleep 0.1
47
    end
48

49
    nil
50
  end
51

52
  def inspect
53
    "#<HistoryManager stack size: #{@contexts.length}>"
54
  end
55

56
  def _contexts
57
    @contexts
58
  end
59

60
  def _debug=(value)
61
    @debug = value
62
  end
63

64
  def _close
65
    event = { type: :close }
66
    @write_queue << event
67
    @remaining_work << event
68
  end
69

70
  private
71

72
  def debug?
73
    @debug
74
  end
75

76
  # A wrapper around mapping the input library to its history; this way we can mock the return value of this method.
77
  def map_library_to_history(input_library)
78
    case input_library
79
    when :readline
80
      ::Readline::HISTORY
81
    when :reline
82
      ::Reline::HISTORY
83
    else
84
      $stderr.puts("Unknown input library: #{input_library}") if debug?
85
      []
86
    end
87
  end
88

89
  def clear_library(input_library)
90
    case input_library
91
    when :readline
92
      clear_readline
93
    when :reline
94
      clear_reline
95
    else
96
      $stderr.puts("Unknown input library: #{input_library}") if debug?
97
    end
98
  end
99

100
  def push_context(history_file: nil, name: nil, input_library: nil)
101
    $stderr.puts("Push context before\n#{JSON.pretty_generate(_contexts)}") if debug?
102
    new_context = { history_file: history_file, name: name, input_library: input_library || :readline }
103

104
    switch_context(new_context, @contexts.last)
105
    @contexts.push(new_context)
106
    $stderr.puts("Push context after\n#{JSON.pretty_generate(_contexts)}") if debug?
107

108
    nil
109
  end
110

111
  def pop_context
112
    $stderr.puts("Pop context before\n#{JSON.pretty_generate(_contexts)}") if debug?
113
    return if @contexts.empty?
114

115
    old_context = @contexts.pop
116
    $stderr.puts("Pop context after\n#{JSON.pretty_generate(_contexts)}") if debug?
117
    switch_context(@contexts.last, old_context)
118

119
    nil
120
  end
121

122
  def readline_available?
123
    defined?(::Readline)
124
  end
125

126
  def reline_available?
127
    begin
128
      require 'reline'
129
      defined?(::Reline)
130
    rescue ::LoadError => _e
131
      false
132
    end
133
  end
134

135
  def clear_readline
136
    return unless readline_available?
137

138
    ::Readline::HISTORY.length.times { ::Readline::HISTORY.pop }
139
  end
140

141
  def clear_reline
142
    return unless reline_available?
143

144
    ::Reline::HISTORY.length.times { ::Reline::HISTORY.pop }
145
  end
146

147
  def load_history_file(context)
148
    history_file = context[:history_file]
149
    history = map_library_to_history(context[:input_library])
150

151
    begin
152
      File.open(history_file, 'rb') do |f|
153
        clear_library(context[:input_library])
154
        f.each_line(chomp: true) do |line|
155
          if context[:input_library] == :reline && history.last&.end_with?("\\")
156
            history.last.delete_suffix!("\\")
157
            history.last << "\n" << line
158
          else
159
            history << line
160
          end
161
        end
162
      end
163
    rescue Errno::EACCES, Errno::ENOENT => e
164
      elog "Failed to open history file: #{history_file} with error: #{e}"
165
    end
166
  end
167

168
  def store_history_file(context)
169
    history_file = context[:history_file]
170
    history = map_library_to_history(context[:input_library])
171

172
    history_diff = history.length < MAX_HISTORY ? history.length : MAX_HISTORY
173

174
    cmds = []
175
    history_diff.times do
176
      entry = history.pop
177
      cmds << entry.scrub.split("\n").join("\\\n")
178
    end
179

180
    write_history_file(history_file, cmds.reverse)
181
  end
182

183
  def switch_context(new_context, old_context=nil)
184
    if old_context && old_context[:history_file]
185
      store_history_file(old_context)
186
    end
187

188
    if new_context && new_context[:history_file]
189
      load_history_file(new_context)
190
    else
191
      clear_readline
192
      clear_reline
193
    end
194
  rescue SignalException => _e
195
    clear_readline
196
    clear_reline
197
  end
198

199
  def write_history_file(history_file, cmds)
200
    write_queue_ref = @write_queue
201
    remaining_work_ref = @remaining_work
202

203
    @write_thread ||= Rex::ThreadFactory.spawn("HistoryManagerWriter", false) do
204
      while (event = write_queue_ref.pop)
205
        begin
206
          break if event[:type] == :close
207

208
          history_file = event[:history_file]
209
          cmds = event[:cmds]
210

211
          File.open(history_file, 'wb+') do |f|
212
            f.puts(cmds)
213
          end
214

215
        rescue => e
216
          elog(e)
217
        ensure
218
          remaining_work_ref.pop
219
        end
220
      end
221
    end
222

223
    event = { type: :write, history_file: history_file, cmds: cmds }
224
    @write_queue << event
225
    @remaining_work << event
226
  end
227
end
228

229
end
230
end
231
end
232
end
233

234