From 6d6b33968f3b7455d062c9a42a240282d9010e48 Mon Sep 17 00:00:00 2001
From: Aaron Soto <aaron_soto@rapid7.com>
Date: Tue, 30 Apr 2019 10:07:20 -0500
Subject: [PATCH] Record UUID upon pingback_reverse_tcp generation and callback

---
 lib/msf/base/sessions/pingback.rb         | 19 ++++++++++++-------
 lib/msf/base/sessions/pingback_options.rb | 16 +++++++++++++---
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/lib/msf/base/sessions/pingback.rb b/lib/msf/base/sessions/pingback.rb
index f0c7fbac55..62337b8222 100644
--- a/lib/msf/base/sessions/pingback.rb
+++ b/lib/msf/base/sessions/pingback.rb
@@ -37,21 +37,26 @@ class Pingback
     if uuid_raw
       uuid_string = uuid_raw.each_byte.map { |b| "%02x" % b.to_i() }.join
       puts("Incoming Pingback_UUID = |" + uuid_string + "|")
-      #asoto-r7, check the database for the entery here!
+
+      res = Mdm::Payload.find_by uuid: uuid_string
+      require 'pry'; binding.pry
+
       begin
-        uuid_original = opts[:datastore]['PingbackUUID'].to_s
-        puts("Original UUID =          |" + uuid_original.gsub("-", "") + "|")
-        if uuid_original.gsub("-", "") == uuid_string
-          puts("UUIDs Match!")
+        if res.nil?
+          puts("Provided UUID (#{uuid_string}) was not found in database!")
+          #TODO: Abort, somehow?
         else
-          puts("UUIDs DO NOT Match!")
+          puts("UUID identified (#{uuid_string})")
         end
       rescue => e
-        puts("Can't get original UUID")
+        #TODO: Can we have a more specific exception handler?
+        #       Test: what if we send no bytes back?  What if we send less than 16 bytes?  Or more than?
+        puts "Can't get original UUID"
         puts "Exception Class: #{ e.class.name }"
         puts "Exception Message: #{ e.message }"
         puts "Exception Backtrace: #{ e.backtrace }"
       end
+      conn.close
     end
     nil
   end
diff --git a/lib/msf/base/sessions/pingback_options.rb b/lib/msf/base/sessions/pingback_options.rb
index 7705fb6952..53a4d4fb86 100644
--- a/lib/msf/base/sessions/pingback_options.rb
+++ b/lib/msf/base/sessions/pingback_options.rb
@@ -47,8 +47,9 @@ module PingbackOptions
     end
 
   end
-    def generate_pingback_uuid
-    puts("generate_pingback_uuid")
+
+  def generate_pingback_uuid
+    puts("Msf::Sessions::PingbackOptions.generate_pingback_uuid")
     conf = {}
     if datastore['PingbackUUID'].to_s.length > 0
       #
@@ -59,7 +60,16 @@ module PingbackOptions
     conf[:pingback_store] = datastore['PingbackUUIDDatabase']
     pingback = Msf::Payload::Pingback.new(conf)
     datastore['PingbackUUID'] ||= pingback.uuid
-    #asoto-r7, this is where we write the UUID to the database.
+
+    require 'pry'; binding.pry
+
+    vprint_status("Writing UUID #{datastore['PingbackUUID']} to database...")
+    Mdm::Payload.create!(name: datastore['PayloadUUIDName'],
+                         uuid: datastore['PingbackUUID'].gsub('-',''),
+                         description: 'pingback',
+                         platform: platform.platforms.first.realname.downcase,
+                         workspace: framework.db.workspace)
+
     pingback.uuid
   end