From 17691b9ff0f1ab3f43205d198f9764d79a189179 Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Fri, 3 Apr 2026 13:02:11 +0000
Subject: [PATCH] feat(ansible): support custom reboot commands

---
 modules.go          | 13 ++++++++++++-
 modules_adv_test.go | 19 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/modules.go b/modules.go
index c76300b..7e31f73 100644
--- a/modules.go
+++ b/modules.go
@@ -3329,6 +3329,7 @@ func (e *Executor) moduleReboot(ctx context.Context, client sshExecutorClient, a
 	postRebootDelay := getIntArg(args, "post_reboot_delay", 0)
 	rebootTimeout := getIntArg(args, "reboot_timeout", 600)
 	testCommand := getStringArg(args, "test_command", "whoami")
+	rebootCommand := getStringArg(args, "reboot_command", "")
 
 	msg := getStringArg(args, "msg", "Reboot initiated by Ansible")
 	runReboot := func(cmd string) (*TaskResult, error) {
@@ -3343,7 +3344,17 @@ func (e *Executor) moduleReboot(ctx context.Context, client sshExecutorClient, a
 		return nil, nil
 	}
 
-	if preRebootDelay > 0 {
+	if rebootCommand != "" {
+		if preRebootDelay > 0 {
+			if result, err := runReboot(sprintf("sleep %d && %s", preRebootDelay, rebootCommand)); err != nil || result != nil {
+				return result, err
+			}
+		} else {
+			if result, err := runReboot(rebootCommand); err != nil || result != nil {
+				return result, err
+			}
+		}
+	} else if preRebootDelay > 0 {
 		cmd := sprintf("sleep %d && shutdown -r now '%s' &", preRebootDelay, msg)
 		if result, err := runReboot(cmd); err != nil || result != nil {
 			return result, err
diff --git a/modules_adv_test.go b/modules_adv_test.go
index d8fb1f3..00e285a 100644
--- a/modules_adv_test.go
+++ b/modules_adv_test.go
@@ -2040,6 +2040,25 @@ func TestModulesAdv_ModuleReboot_Good_WaitsForTestCommand(t *testing.T) {
 	assert.True(t, mock.hasExecuted(`whoami`))
 }
 
+func TestModulesAdv_ModuleReboot_Good_CustomRebootCommand(t *testing.T) {
+	e, mock := newTestExecutorWithMock("host1")
+	mock.expectCommand(`sleep 1 && /sbin/reboot`, "", "", 0)
+	mock.expectCommand(`whoami`, "root\n", "", 0)
+
+	result, err := e.moduleReboot(context.Background(), mock, map[string]any{
+		"reboot_command":   "/sbin/reboot",
+		"pre_reboot_delay": 1,
+		"reboot_timeout":   5,
+	})
+
+	require.NoError(t, err)
+	assert.True(t, result.Changed)
+	assert.Equal(t, "Reboot initiated", result.Msg)
+	assert.Equal(t, 2, mock.commandCount())
+	assert.True(t, mock.hasExecuted(`sleep 1 && /sbin/reboot`))
+	assert.True(t, mock.hasExecuted(`whoami`))
+}
+
 func TestModulesAdv_ModuleReboot_Bad_TimesOutWaitingForTestCommand(t *testing.T) {
 	e, mock := newTestExecutorWithMock("host1")
 	mock.expectCommand(`shutdown -r now 'Reboot initiated by Ansible' &`, "", "", 0)