proxy.rb in ronin-2.0.1

- old
+ new

@@ -15,10 +15,11 @@
 # You should have received a copy of the GNU General Public License
 # along with Ronin.  If not, see <https://www.gnu.org/licenses/>.
 #
 
 require 'ronin/cli/command'
+require 'ronin/cli/host_and_port'
 require 'ronin/core/cli/logging'
 
 require 'ronin/support/network/tcp/proxy'
 require 'ronin/support/network/ssl/proxy'
 require 'ronin/support/network/tls/proxy'
@@ -74,10 +75,11 @@
       #                                      listen on.
       #     REMOTE_HOST:REMOTE_PORT          The remote server to proxy data to.
       #
       class Proxy < Command
 
+        include HostAndPort
         include Core::CLI::Logging
 
         usage '[PROXY_HOST:]PROXY_PORT UPSTREAM_HOST:UPSTREAM_PORT'
 
         option :tcp, short: '-t',
@@ -113,26 +115,26 @@
                            @rewrite << parse_rewrite_rule(value)
                          end
 
         option :rewrite_client, value: {
                                   type:  String,
-                                  usage: '/REGEXP/:STRING',
+                                  usage: '/REGEXP/:STRING'
                                 },
                                 desc: 'Client rewrite rules' do |value|
                                   @rewrite_client << parse_rewrite_rule(value)
                                 end
 
         option :rewrite_server, value: {
                                   type:  String,
-                                  usage: '/REGEXP/:STRING',
+                                  usage: '/REGEXP/:STRING'
                                 },
                                 desc: 'Server rewrite rules' do |value|
                                   @rewrite_server << parse_rewrite_rule(value)
                                 end
 
         option :ignore, short: '-i',
-                        value: {type:  Regexp},
+                        value: {type: Regexp},
                         desc: 'Ignore rules' do |regexp|
                           @ignore << regexp
                         end
 
         option :ignore_client, value: {type: Regexp},
@@ -144,11 +146,11 @@
                                desc: 'Server ignore rules' do |regexp|
                                  @ignore_server << regexp
                                end
 
         option :close, short: '-C',
-                       value: {type:  Regexp},
+                       value: {type: Regexp},
                        desc: 'Close rules' do |regexp|
                          @close << regexp
                        end
 
         option :close_client, value: {type: Regexp},
@@ -292,18 +294,16 @@
         #
         def run(*args)
           local, upstream = *args
 
           if local.include?(':')
-            proxy_host, proxy_port = local.split(':',2)
-            proxy_port = proxy_port.to_i
+            proxy_host, proxy_port = host_and_port(local)
           else
             proxy_port = local.to_i
           end
 
-          upstream_host, upstream_port = upstream.split(':',2)
-          upstream_port = upstream_port.to_i
+          upstream_host, upstream_port = host_and_port(upstream)
 
           if options[:hexdump]
             @hexdumper = Hexdump::Hexdump.new
           end
 
@@ -428,11 +428,11 @@
           unless (index = string.rindex('/:'))
             raise(OptionParser::InvalidArgument,"invalid rewrite rule: #{string}")
           end
 
           regexp  = Regexp.new(string[1...index])
-          pattern = string[index+2..]
+          pattern = string[(index + 2)..]
 
           return regexp, pattern
         end
 
         #
@@ -463,10 +463,10 @@
         #   The address of the connection.
         #
         def address(connection)
           case connection
           when Array
-            socket, (host, port) = connection
+            _socket, (host, port) = connection
 
             "#{host}:#{port}"
           when TCPSocket, UDPSocket
             addrinfo = connection.peeraddr
 