Skip to content

Commit 3b432c0

Browse files
h3zh1h3zh1
committed
feat(cli): add daemon mode for headless runtimes
1 parent 16b9f88 commit 3b432c0

5 files changed

Lines changed: 118 additions & 5 deletions

File tree

‎client/cmd/cli/root.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,15 @@ func rootCmd(con *core.Console) (*cobra.Command, error) {
1313
var cmd = &cobra.Command{
1414
Use: "client",
1515
RunE: func(cmd *cobra.Command, args []string) error {
16+
if err := common.ValidateExecutionModeFlags(cmd); err != nil {
17+
return err
18+
}
1619
if err := generic.LoginCmd(cmd, con); err != nil {
1720
return err
1821
}
19-
if common.ShouldStartConsole(cmd) {
22+
if common.ShouldStartRuntime(cmd) {
23+
restoreDaemon := con.WithDaemonExecution(common.ShouldStartDaemon(cmd))
24+
defer restoreDaemon()
2025
return con.Start(command.BindClientsCommands, command.BindImplantCommands)
2126
}
2227
return nil
@@ -28,6 +33,7 @@ func rootCmd(con *core.Console) (*cobra.Command, error) {
2833
cmd.PersistentFlags().String("mcp", "", "enable MCP server with address (e.g., 127.0.0.1:5005)")
2934
// Add --rpc flag
3035
cmd.PersistentFlags().String("rpc", "", "enable local gRPC server with address (e.g., 127.0.0.1:15004)")
36+
cmd.PersistentFlags().Bool("daemon", false, "keep background services alive without entering the interactive console")
3137
bind := command.MakeBind(cmd, con, "golang")
3238
command.BindCommonCommands(bind)
3339
// Setup console runner

‎client/command/client.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,9 @@ func ConsoleRunnerCmd(con *core.Console, cmd *cobra.Command) (pre, post func(cmd
6969
})
7070

7171
pre = func(cmd *cobra.Command, args []string) error {
72+
if err := common.ValidateExecutionModeFlags(cmd); err != nil {
73+
return err
74+
}
7275
if cmd.Use == consts.CommandLogin || cmd.Use == consts.ClientMenu {
7376
return nil
7477
}
@@ -81,7 +84,9 @@ func ConsoleRunnerCmd(con *core.Console, cmd *cobra.Command) (pre, post func(cmd
8184
return nil
8285
}
8386

84-
if shouldStartConsole(cmd) {
87+
if common.ShouldStartRuntime(cmd) {
88+
restoreDaemon := con.WithDaemonExecution(common.ShouldStartDaemon(cmd))
89+
defer restoreDaemon()
8590
return con.Start(BindClientsCommands, BindImplantCommands)
8691
}
8792

‎client/command/client_test.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ func TestShouldStartConsoleRequiresTTYForLogin(t *testing.T) {
1515

1616
cmd := &cobra.Command{Use: consts.CommandLogin}
1717
cmd.Flags().Bool("console", false, "")
18+
cmd.Flags().Bool("daemon", false, "")
1819

1920
if shouldStartConsole(cmd) {
2021
t.Fatal("login should not start an interactive console in non-interactive mode")
@@ -28,6 +29,7 @@ func TestShouldStartConsoleAllowsExplicitConsoleFlag(t *testing.T) {
2829

2930
cmd := &cobra.Command{Use: "version"}
3031
cmd.Flags().Bool("console", false, "")
32+
cmd.Flags().Bool("daemon", false, "")
3133
if err := cmd.Flags().Set("console", "true"); err != nil {
3234
t.Fatalf("failed to set console flag: %v", err)
3335
}
@@ -44,8 +46,48 @@ func TestShouldStartConsoleDoesNotStartRootInNonInteractiveMode(t *testing.T) {
4446

4547
cmd := &cobra.Command{Use: consts.ClientMenu}
4648
cmd.Flags().Bool("console", false, "")
49+
cmd.Flags().Bool("daemon", false, "")
4750

4851
if shouldStartConsole(cmd) {
4952
t.Fatal("root command should not start the console in non-interactive mode")
5053
}
5154
}
55+
56+
func TestShouldStartDaemonAllowsHeadlessRuntime(t *testing.T) {
57+
old := common.StdinIsTerminal
58+
common.StdinIsTerminal = func() bool { return false }
59+
t.Cleanup(func() { common.StdinIsTerminal = old })
60+
61+
cmd := &cobra.Command{Use: consts.CommandLogin}
62+
cmd.Flags().Bool("console", false, "")
63+
cmd.Flags().Bool("daemon", false, "")
64+
if err := cmd.Flags().Set("daemon", "true"); err != nil {
65+
t.Fatalf("failed to set daemon flag: %v", err)
66+
}
67+
68+
if !common.ShouldStartDaemon(cmd) {
69+
t.Fatal("--daemon should enable headless runtime mode")
70+
}
71+
if !common.ShouldStartRuntime(cmd) {
72+
t.Fatal("--daemon should keep runtime alive even without REPL")
73+
}
74+
if common.ShouldSuppressStartupOutput(cmd) {
75+
t.Fatal("--daemon should preserve startup output for diagnostics")
76+
}
77+
}
78+
79+
func TestValidateExecutionModeFlagsRejectsConsoleAndDaemon(t *testing.T) {
80+
cmd := &cobra.Command{Use: consts.CommandLogin}
81+
cmd.Flags().Bool("console", false, "")
82+
cmd.Flags().Bool("daemon", false, "")
83+
if err := cmd.Flags().Set("console", "true"); err != nil {
84+
t.Fatalf("failed to set console flag: %v", err)
85+
}
86+
if err := cmd.Flags().Set("daemon", "true"); err != nil {
87+
t.Fatalf("failed to set daemon flag: %v", err)
88+
}
89+
90+
if err := common.ValidateExecutionModeFlags(cmd); err == nil {
91+
t.Fatal("expected --console and --daemon to conflict")
92+
}
93+
}

‎client/command/common/execution.go‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package common
22

33
import (
4+
"fmt"
5+
46
"github.com/chainreactors/IoM-go/consts"
57
"github.com/spf13/cobra"
68
)
@@ -24,6 +26,33 @@ func ShouldStartConsole(cmd *cobra.Command) bool {
2426
return cmd == cmd.Root() || cmd.Use == consts.CommandLogin
2527
}
2628

29+
func ShouldStartDaemon(cmd *cobra.Command) bool {
30+
if cmd == nil {
31+
return false
32+
}
33+
34+
run, _ := cmd.Flags().GetBool("daemon")
35+
return run
36+
}
37+
38+
func ShouldStartRuntime(cmd *cobra.Command) bool {
39+
return ShouldStartConsole(cmd) || ShouldStartDaemon(cmd)
40+
}
41+
2742
func ShouldSuppressStartupOutput(cmd *cobra.Command) bool {
28-
return !ShouldStartConsole(cmd)
43+
return !ShouldStartRuntime(cmd)
44+
}
45+
46+
func ValidateExecutionModeFlags(cmd *cobra.Command) error {
47+
if cmd == nil {
48+
return nil
49+
}
50+
51+
runConsole, _ := cmd.Flags().GetBool("console")
52+
runDaemon, _ := cmd.Flags().GetBool("daemon")
53+
if runConsole && runDaemon {
54+
return fmt.Errorf("--console and --daemon cannot be used together")
55+
}
56+
57+
return nil
2958
}

‎client/core/console.go‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ type Console struct {
104104
MalManager *plugin.MalManager
105105

106106
forceNonInteractive atomic.Int32
107+
daemonMode atomic.Int32
107108
replActive atomic.Bool
108109
}
109110

@@ -170,9 +171,14 @@ func (c *Console) Start(bindCmds ...BindCmds) error {
170171
}
171172

172173
// Headless mode: stdin is not a terminal (e.g., launched by GUI with /dev/null).
174+
// Daemon mode follows the same runtime path even when a terminal is available.
173175
// Skip readline loop to avoid busy-spin on MakeRaw(ENOTTY), block on signal instead.
174-
if !term.IsTerminal(int(os.Stdin.Fd())) {
175-
logs.Log.Importantf("running in headless mode (no terminal detected), waiting for signal...")
176+
if c.IsDaemonExecution() || !term.IsTerminal(int(os.Stdin.Fd())) {
177+
if c.IsDaemonExecution() {
178+
logs.Log.Importantf("running in daemon mode, waiting for signal...")
179+
} else {
180+
logs.Log.Importantf("running in headless mode (no terminal detected), waiting for signal...")
181+
}
176182
sig := make(chan os.Signal, 1)
177183
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
178184
<-sig
@@ -231,6 +237,31 @@ func (c *Console) WithNonInteractiveExecution(enabled bool) func() {
231237
}
232238
}
233239

240+
func (c *Console) WithDaemonExecution(enabled bool) func() {
241+
if c == nil {
242+
return func() {}
243+
}
244+
245+
prev := c.daemonMode.Load()
246+
if enabled {
247+
c.daemonMode.Store(1)
248+
} else {
249+
c.daemonMode.Store(0)
250+
}
251+
252+
return func() {
253+
c.daemonMode.Store(prev)
254+
}
255+
}
256+
257+
func (c *Console) IsDaemonExecution() bool {
258+
if c == nil {
259+
return false
260+
}
261+
262+
return c.daemonMode.Load() > 0
263+
}
264+
234265
func (c *Console) WithREPLExecution(enabled bool) func() {
235266
if c == nil {
236267
return func() {}

0 commit comments

Comments
 (0)