Skip to content

Commit dabf050

Browse files
author
mrachidi
committed
formating
1 parent 0680be3 commit dabf050

27 files changed

+1433
-1116
lines changed

mcp_simple_tool/server.py

Lines changed: 51 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,17 @@
11
import anyio
22
import click
3-
from datetime import datetime
43
import mcp.types as types
54
from mcp.server.lowlevel import Server
65

76
# Import tools from modules
87
from .tools.cve_lookup import lookup_cve
9-
from .tools.package_vulnerability import check_package_vulnerabilities
10-
from .tools.epss_lookup import get_epss_score
118
from .tools.cvss_calculator import calculate_cvss_score
12-
from .tools.vulnerability_search import search_vulnerabilities
9+
from .tools.epss_lookup import get_epss_score
1310
from .tools.exploit_availability import get_exploit_availability
14-
from .tools.vulnerability_timeline import get_vulnerability_timeline
11+
from .tools.package_vulnerability import check_package_vulnerabilities
1512
from .tools.vex_status import get_vex_status
13+
from .tools.vulnerability_search import search_vulnerabilities
14+
from .tools.vulnerability_timeline import get_vulnerability_timeline
1615

1716

1817
@click.command()
@@ -92,74 +91,81 @@ def main(port: int, transport: str) -> int:
9291
)
9392

9493
@app.call_tool()
95-
async def fetch_tool( # type: ignore[unused-function]
94+
async def fetch_tool( # type: ignore[unused-function]
9695
name: str, arguments: dict
9796
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
9897
if name == "cve_lookup":
9998
if "cve_id" not in arguments:
100-
return [types.TextContent(
101-
type="text",
102-
text="Error: Missing required argument 'cve_id'"
103-
)]
99+
return [
100+
types.TextContent(
101+
type="text", text="Error: Missing required argument 'cve_id'"
102+
)
103+
]
104104
return await lookup_cve(arguments["cve_id"])
105105
elif name == "package_vulnerability_check":
106106
if "package_name" not in arguments:
107-
return [types.TextContent(
108-
type="text",
109-
text="Error: Missing required argument 'package_name'"
110-
)]
107+
return [
108+
types.TextContent(
109+
type="text",
110+
text="Error: Missing required argument 'package_name'",
111+
)
112+
]
111113
version = arguments.get("version") # Optional parameter
112-
return await check_package_vulnerabilities(arguments["package_name"], version)
114+
return await check_package_vulnerabilities(
115+
arguments["package_name"], version
116+
)
113117
elif name == "get_epss_score":
114118
if "cve_id" not in arguments:
115-
return [types.TextContent(
116-
type="text",
117-
text="Error: Missing required argument 'cve_id'"
118-
)]
119+
return [
120+
types.TextContent(
121+
type="text", text="Error: Missing required argument 'cve_id'"
122+
)
123+
]
119124
return await get_epss_score(arguments["cve_id"])
120125
elif name == "calculate_cvss_score":
121126
if "vector" not in arguments:
122-
return [types.TextContent(
123-
type="text",
124-
text="Error: Missing required argument 'vector'"
125-
)]
127+
return [
128+
types.TextContent(
129+
type="text", text="Error: Missing required argument 'vector'"
130+
)
131+
]
126132
return await calculate_cvss_score(arguments["vector"])
127133
elif name == "search_vulnerabilities":
128134
# All parameters are optional for search
129135
keywords = arguments.get("keywords")
130-
severity = arguments.get("severity")
136+
severity = arguments.get("severity")
131137
date_range = arguments.get("date_range")
132138
return await search_vulnerabilities(keywords, severity, date_range)
133139
elif name == "get_exploit_availability":
134140
if "cve_id" not in arguments:
135-
return [types.TextContent(
136-
type="text",
137-
text="Error: Missing required argument 'cve_id'"
138-
)]
141+
return [
142+
types.TextContent(
143+
type="text", text="Error: Missing required argument 'cve_id'"
144+
)
145+
]
139146
return await get_exploit_availability(arguments["cve_id"])
140147
elif name == "get_vulnerability_timeline":
141148
if "cve_id" not in arguments:
142-
return [types.TextContent(
143-
type="text",
144-
text="Error: Missing required argument 'cve_id'"
145-
)]
149+
return [
150+
types.TextContent(
151+
type="text", text="Error: Missing required argument 'cve_id'"
152+
)
153+
]
146154
return await get_vulnerability_timeline(arguments["cve_id"])
147155
elif name == "get_vex_status":
148156
if "cve_id" not in arguments:
149-
return [types.TextContent(
150-
type="text",
151-
text="Error: Missing required argument 'cve_id'"
152-
)]
157+
return [
158+
types.TextContent(
159+
type="text", text="Error: Missing required argument 'cve_id'"
160+
)
161+
]
153162
product = arguments.get("product") # Optional parameter
154163
return await get_vex_status(arguments["cve_id"], product)
155164
else:
156-
return [types.TextContent(
157-
type="text",
158-
text=f"Error: Unknown tool: {name}"
159-
)]
165+
return [types.TextContent(type="text", text=f"Error: Unknown tool: {name}")]
160166

161167
@app.list_tools()
162-
async def list_tools() -> list[types.Tool]: # type: ignore[unused-function]
168+
async def list_tools() -> list[types.Tool]: # type: ignore[unused-function]
163169
return [
164170
types.Tool(
165171
name="cve_lookup",
@@ -189,7 +195,7 @@ async def list_tools() -> list[types.Tool]: # type: ignore[unused-function]
189195
"version": {
190196
"type": "string",
191197
"description": "Specific version to check (optional). If not provided, checks all known versions.",
192-
}
198+
},
193199
},
194200
},
195201
),
@@ -238,7 +244,7 @@ async def list_tools() -> list[types.Tool]: # type: ignore[unused-function]
238244
"date_range": {
239245
"type": "string",
240246
"description": "Date range filter. Use predefined ranges (30d, 90d, 1y, 2y) or custom format YYYY-MM-DD,YYYY-MM-DD",
241-
}
247+
},
242248
},
243249
},
244250
),
@@ -284,10 +290,10 @@ async def list_tools() -> list[types.Tool]: # type: ignore[unused-function]
284290
"product": {
285291
"type": "string",
286292
"description": "Product name or identifier to check VEX status for (optional). Examples: 'Windows 11', 'RHEL 8', 'Ubuntu 22.04', 'Apache HTTP Server'",
287-
}
293+
},
288294
},
289295
},
290-
)
296+
),
291297
]
292298

293299
if transport == "sse":

mcp_simple_tool/tools/__init__.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,21 @@
55
"""
66

77
from .cve_lookup import lookup_cve
8-
from .package_vulnerability import check_package_vulnerabilities
9-
from .epss_lookup import get_epss_score
108
from .cvss_calculator import calculate_cvss_score
11-
from .vulnerability_search import search_vulnerabilities
9+
from .epss_lookup import get_epss_score
1210
from .exploit_availability import get_exploit_availability
13-
from .vulnerability_timeline import get_vulnerability_timeline
11+
from .package_vulnerability import check_package_vulnerabilities
1412
from .vex_status import get_vex_status
13+
from .vulnerability_search import search_vulnerabilities
14+
from .vulnerability_timeline import get_vulnerability_timeline
1515

1616
__all__ = [
1717
"lookup_cve",
18-
"check_package_vulnerabilities",
18+
"check_package_vulnerabilities",
1919
"get_epss_score",
2020
"calculate_cvss_score",
2121
"search_vulnerabilities",
2222
"get_exploit_availability",
2323
"get_vulnerability_timeline",
24-
"get_vex_status"
25-
]
24+
"get_vex_status",
25+
]

0 commit comments

Comments
 (0)