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
require 'rex/post/sql/ui/console/command_dispatcher'
2
require 'rex/post/sql/ui/console/interactive_sql_client'
3
require 'rex/post/session_compatible_modules'
4

5
module Rex
6
  module Post
7
    module Sql
8
      module Ui
9

10
        #
11
        # Base console class for Generic SQL consoles
12
        #
13
        module Console
14

15
          include Rex::Ui::Text::DispatcherShell
16
          include Rex::Post::SessionCompatibleModules
17

18
          # Called when someone wants to interact with an SQL client.  It's
19
          # assumed that init_ui has been called prior.
20
          #
21
          # @param [Proc] block
22
          # @return [Integer]
23
          def interact(&block)
24
            # Run queued commands
25
            commands.delete_if do |ent|
26
              run_single(ent)
27
              true
28
            end
29

30
            # Run the interactive loop
31
            run do |line|
32
              # Run the command
33
              run_single(line)
34

35
              # If a block was supplied, call it, otherwise return false
36
              if block
37
                block.call
38
              else
39
                false
40
              end
41
            end
42
          end
43

44
          # Queues a command to be run when the interactive loop is entered.
45
          #
46
          # @param [Object] cmd
47
          # @return [Object]
48
          def queue_cmd(cmd)
49
            self.commands << cmd
50
          end
51

52
          # Runs the specified command wrapper in something to catch meterpreter
53
          # exceptions.
54
          #
55
          # @param [Object] dispatcher
56
          # @param [Object] method
57
          # @param [Object] arguments
58
          # @return [FalseClass]
59
          def run_command(dispatcher, method, arguments)
60
            begin
61
              super
62
            rescue ::Timeout::Error
63
              log_error('Operation timed out.')
64
            rescue ::Rex::InvalidDestination => e
65
              log_error(e.message)
66
            rescue ::Errno::EPIPE, ::OpenSSL::SSL::SSLError, ::IOError
67
              self.session.kill
68
            rescue ::StandardError => e
69
              log_error("Error running command #{method}: #{e.class} #{e}")
70
              elog(e)
71
            end
72
          end
73

74
          # @param [Hash] opts
75
          # @return [String]
76
          def help_to_s(opts = {})
77
            super + format_session_compatible_modules
78
          end
79

80
          #
81
          # Notification to display when initially interacting with the client via the query_interactive command
82
          #
83
          # @return [String]
84
          def interact_with_client_notification
85
            print_status("Starting interactive SQL shell for #{sql_prompt}")
86
            print_status('SQL commands ending with ; will be executed on the remote server. Use the %grnexit%clr command to exit.')
87
            print_line
88
          end
89

90
          #
91
          # Create prompt via client and session data
92
          #
93
          # @return [String]
94
          def sql_prompt
95
            "#{session.type} @ #{client.peerinfo}#{current_database.blank? ? '' : " (#{current_database})"}"
96
          end
97

98
          #
99
          # Interacts with the supplied client.
100
          #
101
          def interact_with_client(client_dispatcher: nil)
102
            return unless client_dispatcher
103

104
            interact_with_client_notification
105
            client.extend(InteractiveSqlClient) unless client.is_a?(InteractiveSqlClient)
106
            client.on_command_proc = self.on_command_proc if self.on_command_proc && client.respond_to?(:on_command_proc)
107
            client.on_print_proc   = self.on_print_proc if self.on_print_proc && client.respond_to?(:on_print_proc)
108
            client.on_log_proc = method(:log_output) if self.respond_to?(:log_output, true) && client.respond_to?(:on_log_proc)
109
            client.client_dispatcher = client_dispatcher
110

111
            client.interact(input, output)
112
            client.reset_ui
113
          end
114

115
          # @param [Object] val
116
          # @return [String]
117
          def format_prompt(val)
118
            substitute_colors("%und#{sql_prompt}%clr > ", true)
119
          end
120

121
          #
122
          # Log that an error occurred.
123
          #
124
          def log_error(msg)
125
            print_error(msg)
126

127
            elog(msg, session.type)
128

129
            dlog("Call stack:\n#{$ERROR_POSITION.join("\n")}", session.type)
130
          end
131

132
          def current_database
133
            client.current_database
134
          end
135
        end
136
      end
137
    end
138
  end
139
end
140

141