Skip to content

Commit e66c568

Browse files
M09Icclaude
M09Ic
and
claude
committed
feat(module): consolidate module commands under parent and add unload
Refactor flat module commands (list_module, load_module, refresh_module, clear) into a unified `module` parent with subcommands: list, load, unload, refresh, clear. Add `module unload` to unload a plugin bundle by name. Add bundle_map to Modules proto for module-to-bundle mapping in `module list` output. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 4bb5063 commit e66c568

5 files changed

Lines changed: 163 additions & 39 deletions

File tree

‎client/command/modules/commands.go‎

Lines changed: 97 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -16,76 +16,84 @@ import (
1616
)
1717

1818
func Commands(con *core.Console) []*cobra.Command {
19-
listModuleCmd := &cobra.Command{
20-
Use: consts.ModuleListModule,
19+
moduleCmd := &cobra.Command{
20+
Use: consts.CommandModule,
21+
Short: "Module management",
22+
}
23+
24+
listCmd := &cobra.Command{
25+
Use: "list",
2126
Short: "List modules",
22-
// Long: help.FormatLongHelp(consts.ModuleListModule),
2327
RunE: func(cmd *cobra.Command, args []string) error {
2428
return ListModulesCmd(cmd, con)
2529
},
2630
}
2731

28-
loadModuleCmd := &cobra.Command{
29-
Use: consts.ModuleLoadModule + " [module_file]",
32+
loadCmd := &cobra.Command{
33+
Use: "load [module_file]",
3034
Short: "Load module",
31-
// Long: help.FormatLongHelp(consts.ModuleLoadModule),duan
3235
RunE: func(cmd *cobra.Command, args []string) error {
3336
return LoadModuleCmd(cmd, con)
3437
},
3538
Example: `load module from malefic-modules
36-
before loading, you can list the current modules:
39+
before loading, you can list the current modules:
3740
~~~
38-
execute_addon,exec ...
41+
module list
3942
~~~
4043
then you can load module
4144
~~~
42-
load_module --path <module_file.dll>
45+
module load --path <module_file.dll>
4346
~~~
44-
you can see more modules loaded by list_module
47+
you can see more modules loaded by module list
4548
~~~
4649
execute_addon,clear,ps,powershell...
4750
~~~
48-
`}
49-
50-
common.BindFlag(loadModuleCmd, func(f *pflag.FlagSet) {
51+
`,
52+
}
53+
common.BindFlag(loadCmd, func(f *pflag.FlagSet) {
5154
f.String("path", "", "module path")
5255
f.String("modules", "", "modules list,eg: basic,extend")
5356
f.StringP("bundle", "", "", "bundle name")
5457
f.String("3rd", "", "build 3rd-party modules")
5558
f.String("artifact", "", "exist module artifact")
5659
})
57-
common.BindFlagCompletions(loadModuleCmd, func(comp carapace.ActionMap) {
60+
common.BindFlagCompletions(loadCmd, func(comp carapace.ActionMap) {
5861
comp["path"] = carapace.ActionFiles()
5962
comp["modules"] = common.ModulesCompleter()
6063
comp["artifact"] = common.ModuleArtifactsCompleter(con)
6164
})
62-
common.BindArgCompletions(loadModuleCmd, nil,
65+
common.BindArgCompletions(loadCmd, nil,
6366
carapace.ActionFiles().Usage("path to the module file"))
6467

65-
refreshModuleCmd := &cobra.Command{
66-
Use: consts.ModuleRefreshModule,
68+
unloadCmd := &cobra.Command{
69+
Use: "unload [bundle_name]",
70+
Short: "Unload a module bundle by name",
71+
Args: cobra.ExactArgs(1),
72+
RunE: func(cmd *cobra.Command, args []string) error {
73+
return UnloadModuleCmd(cmd, con)
74+
},
75+
}
76+
common.BindArgCompletions(unloadCmd, nil,
77+
common.SessionModuleCompleter(con).Usage("bundle name to unload"))
78+
79+
refreshCmd := &cobra.Command{
80+
Use: "refresh",
6781
Short: "Refresh module",
68-
// Long: help.FormatLongHelp(consts.ModuleRefreshModule),
6982
RunE: func(cmd *cobra.Command, args []string) error {
7083
return RefreshModuleCmd(cmd, con)
7184
},
7285
}
7386

7487
clearCmd := &cobra.Command{
75-
Use: consts.ModuleClear,
76-
Short: "Clear modules",
77-
// Long: help.FormatLongHelp(consts.ModuleClear),
88+
Use: "clear",
89+
Short: "Clear all modules",
7890
RunE: func(cmd *cobra.Command, args []string) error {
7991
return ClearCmd(cmd, con)
8092
},
8193
}
8294

83-
return []*cobra.Command{
84-
listModuleCmd,
85-
loadModuleCmd,
86-
refreshModuleCmd,
87-
clearCmd,
88-
}
95+
moduleCmd.AddCommand(listCmd, loadCmd, unloadCmd, refreshCmd, clearCmd)
96+
return []*cobra.Command{moduleCmd}
8997
}
9098

9199
func Register(con *core.Console) {
@@ -108,17 +116,21 @@ func Register(con *core.Console) {
108116
var rowEntries []table.Row
109117
var row table.Row
110118
tableModel := tui.NewTable([]table.Column{
111-
table.NewFlexColumn("Module", "Module", 1),
112-
table.NewFlexColumn("Help", "Help", 2),
119+
table.NewFlexColumn("Module", "Module", 2),
120+
table.NewFlexColumn("Bundle", "Bundle", 1),
121+
table.NewFlexColumn("Help", "Help", 3),
113122
}, true)
123+
bundleMap := modules.GetBundleMap()
114124
for _, module := range modules.GetModules() {
115125
var short string
116126
if cmd := con.CMDs[module]; cmd != nil {
117127
short = cmd.Short
118128
}
129+
bundle := bundleMap[module]
119130
row = table.NewRow(
120131
table.RowData{
121132
"Module": module,
133+
"Bundle": bundle,
122134
"Help": short,
123135
})
124136
rowEntries = append(rowEntries, row)
@@ -152,6 +164,61 @@ func Register(con *core.Console) {
152164
},
153165
[]string{"task"})
154166

167+
con.RegisterImplantFunc(
168+
consts.ModuleUnloadModule,
169+
unloadModule,
170+
"",
171+
nil,
172+
func(ctx *clientpb.TaskContext) (interface{}, error) {
173+
resp := ctx.Spite.GetModules()
174+
ctx.Session.Modules = resp.Modules
175+
con.RefreshCmd(con.AddSession(ctx.Session))
176+
return resp.Modules, nil
177+
},
178+
func(content *clientpb.TaskContext) (string, error) {
179+
modules := content.Spite.GetModules()
180+
remaining := modules.GetModules()
181+
if len(remaining) == 0 {
182+
return "All modules unloaded.", nil
183+
}
184+
185+
var rowEntries []table.Row
186+
var row table.Row
187+
tableModel := tui.NewTable([]table.Column{
188+
table.NewFlexColumn("Module", "Module", 2),
189+
table.NewFlexColumn("Bundle", "Bundle", 1),
190+
table.NewFlexColumn("Help", "Help", 3),
191+
}, true)
192+
bundleMap := modules.GetBundleMap()
193+
for _, module := range remaining {
194+
var short string
195+
if cmd := con.CMDs[module]; cmd != nil {
196+
short = cmd.Short
197+
}
198+
bundle := bundleMap[module]
199+
row = table.NewRow(
200+
table.RowData{
201+
"Module": module,
202+
"Bundle": bundle,
203+
"Help": short,
204+
})
205+
rowEntries = append(rowEntries, row)
206+
}
207+
tableModel.SetMultiline()
208+
tableModel.SetRows(rowEntries)
209+
return "Unloaded successfully. Remaining modules:\n" + tableModel.View(), nil
210+
})
211+
212+
con.AddCommandFuncHelper(
213+
consts.ModuleUnloadModule,
214+
consts.ModuleUnloadModule,
215+
consts.ModuleUnloadModule+"(active(),\"bundle_name\")",
216+
[]string{
217+
"session: special session",
218+
"bundle: bundle name to unload",
219+
},
220+
[]string{"task"})
221+
155222
con.RegisterImplantFunc(
156223
consts.ModuleRefreshModule,
157224
refreshModule,
@@ -173,7 +240,7 @@ func Register(con *core.Console) {
173240
},
174241
[]string{"task"})
175242

176-
//clear
243+
// clear
177244
con.RegisterImplantFunc(
178245
consts.ModuleClear,
179246
clearAll,

‎client/command/modules/modules_test.go‎

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ func TestLoadModuleFromPath(t *testing.T) {
2020
t.Fatalf("WriteFile failed: %v", err)
2121
}
2222

23-
if err := h.Execute(consts.ModuleLoadModule, "--path", path); err != nil {
23+
if err := h.Execute(consts.CommandModule, "load", "--path", path); err != nil {
2424
t.Fatalf("Execute failed: %v", err)
2525
}
2626

@@ -45,7 +45,7 @@ func TestLoadModuleFromArtifactDownloadsThenLoads(t *testing.T) {
4545
}, nil
4646
})
4747

48-
if err := h.Execute(consts.ModuleLoadModule, "--artifact", "artifact-module.dll"); err != nil {
48+
if err := h.Execute(consts.CommandModule, "load", "--artifact", "artifact-module.dll"); err != nil {
4949
t.Fatalf("Execute failed: %v", err)
5050
}
5151

@@ -87,7 +87,7 @@ func TestLoadModuleBuildUsesSelectedModules(t *testing.T) {
8787
}, nil
8888
})
8989

90-
if err := h.Execute(consts.ModuleLoadModule, "--modules", "nano, execute_dll"); err != nil {
90+
if err := h.Execute(consts.CommandModule, "load", "--modules", "nano, execute_dll"); err != nil {
9191
t.Fatalf("Execute failed: %v", err)
9292
}
9393

@@ -139,7 +139,7 @@ func TestLoadModuleBuildUsesSelectedModules(t *testing.T) {
139139
func TestLoadModuleBuildUsesThirdPartySelection(t *testing.T) {
140140
h := testsupport.NewHarness(t)
141141

142-
if err := h.Execute(consts.ModuleLoadModule, "--3rd", "rem"); err != nil {
142+
if err := h.Execute(consts.CommandModule, "load", "--3rd", "rem"); err != nil {
143143
t.Fatalf("Execute failed: %v", err)
144144
}
145145

@@ -174,7 +174,7 @@ func TestLoadModuleBuildErrorsPropagate(t *testing.T) {
174174
return nil, context.DeadlineExceeded
175175
})
176176

177-
err := h.Execute(consts.ModuleLoadModule, "--modules", "nano")
177+
err := h.Execute(consts.CommandModule, "load", "--modules", "nano")
178178
if err == nil || err != context.DeadlineExceeded {
179179
t.Fatalf("Execute error = %v, want %v", err, context.DeadlineExceeded)
180180
}
@@ -192,7 +192,7 @@ func TestLoadModuleBuildErrorsPropagate(t *testing.T) {
192192
func TestLoadModuleRejectsMutuallyExclusiveSelectors(t *testing.T) {
193193
h := testsupport.NewHarness(t)
194194

195-
err := h.Execute(consts.ModuleLoadModule, "--modules", "nano", "--3rd", "rem")
195+
err := h.Execute(consts.CommandModule, "load", "--modules", "nano", "--3rd", "rem")
196196
if err == nil {
197197
t.Fatal("expected mutually exclusive selector error")
198198
}
@@ -207,7 +207,7 @@ func TestLoadModuleRejectsMutuallyExclusiveSelectors(t *testing.T) {
207207
func TestLoadModuleRejectsMultipleInputSources(t *testing.T) {
208208
h := testsupport.NewHarness(t)
209209

210-
err := h.Execute(consts.ModuleLoadModule, "--artifact", "module.dll", "--modules", "nano")
210+
err := h.Execute(consts.CommandModule, "load", "--artifact", "module.dll", "--modules", "nano")
211211
if err == nil {
212212
t.Fatal("expected multiple input source error")
213213
}
@@ -222,7 +222,7 @@ func TestLoadModuleRejectsMultipleInputSources(t *testing.T) {
222222
func TestLoadModuleRequiresOneSource(t *testing.T) {
223223
h := testsupport.NewHarness(t)
224224

225-
err := h.Execute(consts.ModuleLoadModule)
225+
err := h.Execute(consts.CommandModule, "load")
226226
if err == nil {
227227
t.Fatal("expected missing source error")
228228
}
@@ -234,6 +234,25 @@ func TestLoadModuleRequiresOneSource(t *testing.T) {
234234
testsupport.RequireNoSessionEvents(t, h)
235235
}
236236

237+
func TestUnloadModule(t *testing.T) {
238+
h := testsupport.NewHarness(t)
239+
240+
if err := h.Execute(consts.CommandModule, "unload", "execute_dll"); err != nil {
241+
t.Fatalf("Execute failed: %v", err)
242+
}
243+
244+
req, md := testsupport.MustSingleCall[*implantpb.Request](t, h, "UnloadModule")
245+
if req.Name != consts.ModuleUnloadModule {
246+
t.Fatalf("request name = %q, want %q", req.Name, consts.ModuleUnloadModule)
247+
}
248+
if req.Input != "execute_dll" {
249+
t.Fatalf("request input = %q, want execute_dll", req.Input)
250+
}
251+
testsupport.RequireSessionID(t, md, h.Session.SessionId)
252+
testsupport.RequireCallee(t, md, consts.CalleeCMD)
253+
assertSingleTaskEvent(t, h, consts.ModuleUnloadModule)
254+
}
255+
237256
func assertSingleTaskEvent(t testing.TB, h *testsupport.Harness, wantType string) {
238257
t.Helper()
239258

‎client/command/modules/unload.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package modules
2+
3+
import (
4+
"github.com/chainreactors/IoM-go/client"
5+
"github.com/chainreactors/IoM-go/consts"
6+
"github.com/chainreactors/IoM-go/proto/client/clientpb"
7+
"github.com/chainreactors/IoM-go/proto/implant/implantpb"
8+
"github.com/chainreactors/IoM-go/proto/services/clientrpc"
9+
"github.com/chainreactors/malice-network/client/core"
10+
"github.com/spf13/cobra"
11+
)
12+
13+
func UnloadModuleCmd(cmd *cobra.Command, con *core.Console) error {
14+
bundleName := cmd.Flags().Args()[0]
15+
session := con.GetInteractive()
16+
task, err := unloadModule(con.Rpc, session, bundleName)
17+
if err != nil {
18+
return err
19+
}
20+
session.Console(task, string(*con.App.Shell().Line()))
21+
return nil
22+
}
23+
24+
func unloadModule(rpc clientrpc.MaliceRPCClient, session *client.Session, bundle string) (*clientpb.Task, error) {
25+
return rpc.UnloadModule(session.Context(), &implantpb.Request{
26+
Name: consts.ModuleUnloadModule,
27+
Input: bundle,
28+
})
29+
}

‎external/IoM-go‎

‎server/rpc/rpc-module.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,15 @@ func (rpc *Server) RefreshModule(ctx context.Context, req *implantpb.Request) (*
5454
})
5555
}
5656

57+
func (rpc *Server) UnloadModule(ctx context.Context, req *implantpb.Request) (*clientpb.Task, error) {
58+
if req == nil {
59+
return nil, types.ErrMissingRequestField
60+
}
61+
return rpc.AssertAndHandleWithSession(ctx, req, consts.ModuleUnloadModule, types.MsgListModule, func(greq *GenericRequest, spite *implantpb.Spite) {
62+
applyModulesResponse(greq.Session, spite, false)
63+
})
64+
}
65+
5766
func (rpc *Server) Clear(ctx context.Context, req *implantpb.Request) (*clientpb.Task, error) {
5867
return rpc.AssertAndHandle(ctx, req, consts.ModuleClear, types.MsgEmpty)
5968
}

0 commit comments

Comments
 (0)