Skip to content

Commit 026f6af

Browse files
psiinonpsiinon
committed
mcp: MCP Import WIP
Signed-off-by: Simon Bennetts <psiinon@gmail.com>
1 parent f80a264 commit 026f6af

24 files changed

Lines changed: 2558 additions & 25 deletions

‎addOns/mcp/src/main/java/org/zaproxy/addon/mcp/ExtensionMcp.java‎

Lines changed: 56 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,11 @@
3030
import org.parosproxy.paros.extension.ExtensionHook;
3131
import org.parosproxy.paros.extension.history.ExtensionHistory;
3232
import org.parosproxy.paros.model.OptionsParam;
33+
import org.parosproxy.paros.view.View;
3334
import org.zaproxy.addon.automation.ExtensionAutomation;
35+
import org.zaproxy.addon.mcp.automation.ImportMcpServerJob;
36+
import org.zaproxy.addon.mcp.automation.McpConfigJob;
37+
import org.zaproxy.addon.mcp.importer.ImportMcpServerDialog;
3438
import org.zaproxy.addon.mcp.prompts.ZapBaselineScanPrompt;
3539
import org.zaproxy.addon.mcp.prompts.ZapFullScanPrompt;
3640
import org.zaproxy.addon.mcp.resources.AlertInstancesResource;
@@ -61,6 +65,7 @@
6165
import org.zaproxy.addon.network.server.Server;
6266
import org.zaproxy.addon.pscan.ExtensionPassiveScan2;
6367
import org.zaproxy.addon.reports.ExtensionReports;
68+
import org.zaproxy.zap.view.ZapMenuItem;
6469

6570
/** The MCP Integration add-on extension. */
6671
public class ExtensionMcp extends ExtensionAdaptor {
@@ -80,6 +85,8 @@ public class ExtensionMcp extends ExtensionAdaptor {
8085
private static final Logger LOGGER = LogManager.getLogger(ExtensionMcp.class);
8186

8287
private Server server;
88+
private ImportMcpServerJob importMcpServerJob;
89+
private McpConfigJob mcpConfigJob;
8390
private McpParam param;
8491
private McpToolRegistry toolRegistry;
8592
private McpResourceRegistry resourceRegistry;
@@ -90,6 +97,7 @@ public class ExtensionMcp extends ExtensionAdaptor {
9097
public ExtensionMcp() {
9198
super(NAME);
9299
setI18nPrefix(PREFIX);
100+
this.setOrder(800);
93101
}
94102

95103
@Override
@@ -109,9 +117,28 @@ public void init() {
109117
public void hook(ExtensionHook extensionHook) {
110118
extensionHook.addOptionsParamSet(param);
111119
extensionHook.addOptionsChangedListener(this::optionsChanged);
120+
extensionHook.addVariant(org.zaproxy.addon.mcp.importer.VariantMcpJsonRpc.class);
121+
122+
ExtensionAutomation extAutomation =
123+
Control.getSingleton().getExtensionLoader().getExtension(ExtensionAutomation.class);
124+
if (extAutomation != null) {
125+
importMcpServerJob = new ImportMcpServerJob();
126+
extAutomation.registerAutomationJob(importMcpServerJob);
127+
mcpConfigJob = new McpConfigJob();
128+
extAutomation.registerAutomationJob(mcpConfigJob);
129+
}
112130

113131
if (hasView()) {
114132
extensionHook.getHookView().addOptionPanel(new McpOptionsPanel());
133+
134+
ZapMenuItem importMenuItem = new ZapMenuItem("mcp.importserver.menu");
135+
importMenuItem.addActionListener(
136+
e -> {
137+
ImportMcpServerDialog dialog =
138+
new ImportMcpServerDialog(View.getSingleton().getMainFrame());
139+
dialog.setVisible(true);
140+
});
141+
extensionHook.getHookMenu().addImportMenuItem(importMenuItem);
115142
}
116143

117144
toolRegistry.registerTool(new ZapVersionTool());
@@ -172,8 +199,8 @@ public McpPromptRegistry getPromptRegistry() {
172199
}
173200

174201
@Override
175-
public void optionsLoaded() {
176-
startServer();
202+
public void start() {
203+
applyServerConfig();
177204
}
178205

179206
private void optionsChanged(OptionsParam optionsParam) {
@@ -184,8 +211,23 @@ private void optionsChanged(OptionsParam optionsParam) {
184211
}
185212
}
186213

214+
public McpParam getMcpParam() {
215+
return param;
216+
}
217+
218+
/** Stops and restarts the MCP server using the current param values. */
219+
public void applyServerConfig() {
220+
if (server != null && (!param.isEnabled() || lastPort != param.getPort())) {
221+
stopServer();
222+
}
223+
224+
if (param.isEnabled()) {
225+
startServer();
226+
}
227+
}
228+
187229
private void startServer() {
188-
if (server != null) {
230+
if (server != null || !param.isEnabled()) {
189231
return;
190232
}
191233

@@ -202,6 +244,7 @@ private void startServer() {
202244
this.getAddOn().getVersion().toString()));
203245
try {
204246
int port = server.start(Server.DEFAULT_ADDRESS, param.getPort());
247+
lastPort = port;
205248
LOGGER.info("MCP HTTP listener started on {}:{}", Server.DEFAULT_ADDRESS, port);
206249
} catch (IOException e) {
207250
LOGGER.warn("Failed to start MCP HTTP listener on port {}", param.getPort(), e);
@@ -229,6 +272,16 @@ public void stop() {
229272
@Override
230273
public void unload() {
231274
stopServer();
275+
ExtensionAutomation extAutomation =
276+
Control.getSingleton().getExtensionLoader().getExtension(ExtensionAutomation.class);
277+
if (extAutomation != null) {
278+
if (importMcpServerJob != null) {
279+
extAutomation.unregisterAutomationJob(importMcpServerJob);
280+
}
281+
if (mcpConfigJob != null) {
282+
extAutomation.unregisterAutomationJob(mcpConfigJob);
283+
}
284+
}
232285
}
233286

234287
@Override

‎addOns/mcp/src/main/java/org/zaproxy/addon/mcp/McpOptionsPanel.java‎

Lines changed: 51 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -40,12 +40,13 @@ public class McpOptionsPanel extends AbstractParamPanel {
4040

4141
private static final long serialVersionUID = 1L;
4242

43+
private JCheckBox enabledCheckBox;
4344
private ZapNumberSpinner portSpinner;
4445
private JCheckBox securityKeyEnabledCheckBox;
4546
private JPasswordField securityKeyField;
4647
private JButton generateKeyButton;
47-
private JCheckBox recordInHistoryCheckBox;
4848
private JCheckBox secureOnlyCheckBox;
49+
private JCheckBox recordInHistoryCheckBox;
4950

5051
public McpOptionsPanel() {
5152
super();
@@ -58,6 +59,12 @@ public McpOptionsPanel() {
5859

5960
int row = 0;
6061

62+
panel.add(
63+
getEnabledCheckBox(),
64+
LayoutHelper.getGBC(
65+
0, row, GridBagConstraints.REMAINDER, 1.0, new Insets(2, 2, 2, 2)));
66+
row++;
67+
6168
JLabel portLabel = new JLabel(Constant.messages.getString("mcp.optionspanel.port.label"));
6269
portLabel.setLabelFor(getPortSpinner());
6370
panel.add(
@@ -111,6 +118,26 @@ public McpOptionsPanel() {
111118
add(panel);
112119
}
113120

121+
private JCheckBox getEnabledCheckBox() {
122+
if (enabledCheckBox == null) {
123+
enabledCheckBox =
124+
new JCheckBox(Constant.messages.getString("mcp.optionspanel.enabled.label"));
125+
enabledCheckBox.addItemListener(
126+
e -> {
127+
boolean enabled = enabledCheckBox.isSelected();
128+
getPortSpinner().setEnabled(enabled);
129+
getSecurityKeyField()
130+
.setEnabled(
131+
enabled && getSecurityKeyEnabledCheckBox().isSelected());
132+
getSecurityKeyEnabledCheckBox().setEnabled(enabled);
133+
getGenerateKeyButton().setEnabled(enabled);
134+
getSecureOnlyCheckBox().setEnabled(enabled);
135+
getRecordInHistoryCheckBox().setEnabled(enabled);
136+
});
137+
}
138+
return enabledCheckBox;
139+
}
140+
114141
private ZapNumberSpinner getPortSpinner() {
115142
if (portSpinner == null) {
116143
portSpinner = new ZapNumberSpinner(1, McpParam.DEFAULT_PORT, 65535);
@@ -137,6 +164,14 @@ private JPasswordField getSecurityKeyField() {
137164
return securityKeyField;
138165
}
139166

167+
private JCheckBox getSecureOnlyCheckBox() {
168+
if (secureOnlyCheckBox == null) {
169+
secureOnlyCheckBox =
170+
new JCheckBox(Constant.messages.getString("mcp.optionspanel.secureonly.label"));
171+
}
172+
return secureOnlyCheckBox;
173+
}
174+
140175
private JCheckBox getRecordInHistoryCheckBox() {
141176
if (recordInHistoryCheckBox == null) {
142177
recordInHistoryCheckBox =
@@ -146,14 +181,6 @@ private JCheckBox getRecordInHistoryCheckBox() {
146181
return recordInHistoryCheckBox;
147182
}
148183

149-
private JCheckBox getSecureOnlyCheckBox() {
150-
if (secureOnlyCheckBox == null) {
151-
secureOnlyCheckBox =
152-
new JCheckBox(Constant.messages.getString("mcp.optionspanel.secureonly.label"));
153-
}
154-
return secureOnlyCheckBox;
155-
}
156-
157184
private JButton getGenerateKeyButton() {
158185
if (generateKeyButton == null) {
159186
generateKeyButton =
@@ -164,7 +191,7 @@ private JButton getGenerateKeyButton() {
164191
e -> {
165192
getSecurityKeyField().setText(McpParam.generateRandomKey());
166193
getSecurityKeyEnabledCheckBox().setSelected(true);
167-
getSecurityKeyField().setEnabled(true);
194+
getSecurityKeyField().setEnabled(getEnabledCheckBox().isSelected());
168195
});
169196
}
170197
return generateKeyButton;
@@ -175,16 +202,26 @@ public void initParam(Object obj) {
175202
OptionsParam options = (OptionsParam) obj;
176203
McpParam param = options.getParamSet(McpParam.class);
177204

205+
boolean enabled = param.isEnabled();
206+
getEnabledCheckBox().setSelected(enabled);
178207
getPortSpinner().setValue(param.getPort());
208+
getPortSpinner().setEnabled(enabled);
179209
getSecurityKeyEnabledCheckBox().setSelected(param.isSecurityKeyEnabled());
210+
getSecurityKeyEnabledCheckBox().setEnabled(enabled);
180211
getSecurityKeyField().setText(param.getSecurityKey());
181-
getSecurityKeyField().setEnabled(param.isSecurityKeyEnabled());
182-
getRecordInHistoryCheckBox().setSelected(param.isRecordInHistory());
212+
getSecurityKeyField().setEnabled(enabled && param.isSecurityKeyEnabled());
213+
getGenerateKeyButton().setEnabled(enabled);
183214
getSecureOnlyCheckBox().setSelected(param.isSecureOnly());
215+
getSecureOnlyCheckBox().setEnabled(enabled);
216+
getRecordInHistoryCheckBox().setSelected(param.isRecordInHistory());
217+
getRecordInHistoryCheckBox().setEnabled(enabled);
184218
}
185219

186220
@Override
187221
public void validateParam(Object obj) throws Exception {
222+
if (!getEnabledCheckBox().isSelected()) {
223+
return;
224+
}
188225
int port = getPortSpinner().getValue();
189226
if (port < 1 || port > 65535) {
190227
throw new IllegalStateException(
@@ -204,11 +241,12 @@ public void saveParam(Object obj) throws Exception {
204241
OptionsParam options = (OptionsParam) obj;
205242
McpParam param = options.getParamSet(McpParam.class);
206243

244+
param.setEnabled(getEnabledCheckBox().isSelected());
207245
param.setPort(getPortSpinner().getValue());
208246
param.setSecurityKeyEnabled(getSecurityKeyEnabledCheckBox().isSelected());
209247
param.setSecurityKey(new String(getSecurityKeyField().getPassword()));
210-
param.setRecordInHistory(getRecordInHistoryCheckBox().isSelected());
211248
param.setSecureOnly(getSecureOnlyCheckBox().isSelected());
249+
param.setRecordInHistory(getRecordInHistoryCheckBox().isSelected());
212250
}
213251

214252
@Override

‎addOns/mcp/src/main/java/org/zaproxy/addon/mcp/McpParam.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ public class McpParam extends VersionedAbstractParam {
2828
private static final String MCP_KEY = "mcp";
2929

3030
private static final String CONFIG_VERSION_KEY = MCP_KEY + VERSION_ATTRIBUTE;
31+
private static final String ENABLED_KEY = MCP_KEY + ".enabled";
3132
private static final String PORT_KEY = MCP_KEY + ".port";
3233
private static final String SECURITY_KEY_ENABLED_KEY = MCP_KEY + ".securityKeyEnabled";
3334
private static final String SECURITY_KEY_KEY = MCP_KEY + ".securityKey";
@@ -44,6 +45,7 @@ public class McpParam extends VersionedAbstractParam {
4445
/** Default port for the MCP HTTP listener. */
4546
public static final int DEFAULT_PORT = 8282;
4647

48+
private boolean enabled = true;
4749
private int port = DEFAULT_PORT;
4850
private boolean securityKeyEnabled = true;
4951
private String securityKey;
@@ -69,6 +71,7 @@ protected void updateConfigsImpl(int fileVersion) {
6971

7072
@Override
7173
protected void parseImpl() {
74+
enabled = getBoolean(ENABLED_KEY, true);
7275
port = getInt(PORT_KEY, DEFAULT_PORT);
7376
securityKeyEnabled = getBoolean(SECURITY_KEY_ENABLED_KEY, true);
7477
securityKey = getString(SECURITY_KEY_KEY, null);
@@ -89,6 +92,15 @@ public static String generateRandomKey() {
8992
return sb.toString();
9093
}
9194

95+
public boolean isEnabled() {
96+
return enabled;
97+
}
98+
99+
public void setEnabled(boolean enabled) {
100+
this.enabled = enabled;
101+
getConfig().setProperty(ENABLED_KEY, enabled);
102+
}
103+
92104
public int getPort() {
93105
return port;
94106
}

0 commit comments

Comments
 (0)