From 74acdf2b3208b59768b552be3be8c18deec8e003 Mon Sep 17 00:00:00 2001
From: dledda-r7 <diego_ledda@rapid7.com>
Date: Wed, 29 Jan 2025 05:50:03 -0500
Subject: [PATCH] feat: persistence mixin draft

---
 lib/msf/core/exploit/local/persistence.rb | 55 ++++++++++++++++++-----
 1 file changed, 43 insertions(+), 12 deletions(-)

diff --git a/lib/msf/core/exploit/local/persistence.rb b/lib/msf/core/exploit/local/persistence.rb
index 0bfcf411be84..2a4267bdceeb 100644
--- a/lib/msf/core/exploit/local/persistence.rb
+++ b/lib/msf/core/exploit/local/persistence.rb
@@ -1,30 +1,61 @@
 # -*- coding: binary -*-
 
 module Msf
-  # This module provides methods for persisting on a target system. Mainly initialization
-  # options.
   module Exploit::Local::Persistence
     def initialize(info = {})
+      @persistence_service = Rex::Sync::Event.new(auto_reset=false)
       super(
         update_info(
           info,
-          'DefaultOptions' => {
-            # leaving this commented out, we don't want a wfs delay so that the module
-            # will run forever.
-            # 'WfsDelay' => 25 * 60 * 60, # 25hrs
-            'AllowNoCleanup' => true # don't delete our persistence after we get a shell
-          },
+          'DefaultOptions' => {},
           # https://github.com/rapid7/metasploit-framework/pull/19676#discussion_r1907594308
-          'Stance' => Msf::Exploit::Stance::Passive
-          # 'Passive' => true # XXX when set, ignores wfsdelay and immediately exists after last command
+          'Stance' => Msf::Exploit::Stance::Passive,
+          'Passive' => true,
+          'Actions' => [
+            [ 'INSTALL', { 'Description' => 'Install the persistence' } ],
+            [ 'CLEANUP', { 'Description' => 'Cleanup the persistence' } ]
+          ],
+          'DefaultAction' => 'INSTALL'
         )
       )
 
       register_advanced_options(
         [
-          OptString.new('WritableDir', [true, 'A directory where we can write files', ''])
+          OptString.new('WritableDir', [true, 'A directory where we can write files', '/tmp/']),
+          OptBool.new('CleanUpPersistence', [true, 'Remove the installed persistence at the end of the module', false])
         ]
       )
     end
+
+    def exploit
+      
+      case action.name.upcase
+      when 'INSTALL'
+        run_as_background = !datastore['DisablePayloadHandler']
+        print_warning('Payload handler is disabled, the persistence will be installed only.') unless run_as_background
+
+        # Call the install_persistence function
+        # must be declared inside the persistence module
+        install_persistence
+
+        @persistence_service.wait if run_as_background
+
+        cleanup_persistence if datastore['CleanUpPersistence']
+
+      when 'CLEANUP'
+
+        # call cleanup_persistence
+        # must be declared inside the persistence module
+        cleanup_persistence
+      end
+    end
+
+    def install_persistence
+      # to be overloaded by the module
+    end
+
+    def cleanup
+      # this is done by the action
+    end
   end
-end
+end
\ No newline at end of file