Implement tracking

lib/msf/core/payload/uuid/options.rb

@@ -17,6 +17,8 @@ module Msf::Payload::UUID::Options
       [
         Msf::OptString.new('PayloadUUIDSeed', [ false, 'A string to use when generating the payload UUID (deterministic)']),
         Msf::OptString.new('PayloadUUIDRaw', [ false, 'A hex string representing the raw 8-byte PUID value for the UUID']),
+        Msf::OptString.new('PayloadUUIDName', [ false, 'A human-friendly name to reference this unique payload (requires tracking)']),
+        Msf::OptBool.new('PayloadUUIDTracking', [ true, 'Whether or not to automatically register generated UUIDs', true]),
       ], self.class)
   end
 
@@ -31,7 +33,7 @@ module Msf::Payload::UUID::Options
   def generate_uri_uuid_mode(mode,len=nil)
     sum = uri_checksum_lookup(mode)
 
-    # The URI length may not have room for an embedded checksum
+    # The URI length may not have room for an embedded UUID
     if len && len < URI_CHECKSUM_UUID_MIN_LEN
       # Throw an error if the user set a seed, but there is no room for it
       if datastore['PayloadUUIDSeed'].to_s.length > 0 ||datastore['PayloadUUIDRaw'].to_s.length > 0
@@ -66,7 +68,34 @@ module Msf::Payload::UUID::Options
       conf[:puid] = puid_raw
     end
 
-    Msf::Payload::UUID.new(conf)
+    # Generate the UUID object
+    uuid = Msf::Payload::UUID.new(conf)
+    record_payload_uuid(uuid)
+
+    uuid
+  end
+
+  # Store a UUID in the JSON database if tracking is enabled
+  def record_payload_uuid(uuid, info={})
+    return unless datastore['PayloadUUIDTracking']
+
+    uuid_info = info.merge({
+      arch: uuid.arch,
+      platform: uuid.platform,
+      timestamp: uuid.timestamp,
+      payload: self.fullname,
+      datastore: self.datastore
+    })
+
+    if datastore['PayloadUUIDSeed'].to_s.length > 0
+      uuid_info[:seed] = datastore['PayloadUUIDSeed']
+    end
+
+    if datastore['PayloadUUIDName'].to_s.length > 0
+      uuid_info[:name] = datastore['PayloadUUIDName']
+    end
+
+    framework.uuid_db[uuid.puid_hex] = uuid_info
   end
 
 end