From f5d6994578fc0d83c3877209917fb6a863ce2314 Mon Sep 17 00:00:00 2001 From: pinksawtooth Date: Thu, 3 Jul 2025 00:08:49 +0900 Subject: [PATCH 1/2] fix Hunting tools --- server/gti/gti_mcp/tools/intelligence.py | 43 +++++++++--------------- 1 file changed, 16 insertions(+), 27 deletions(-) diff --git a/server/gti/gti_mcp/tools/intelligence.py b/server/gti/gti_mcp/tools/intelligence.py index 575617fc..7850bcba 100644 --- a/server/gti/gti_mcp/tools/intelligence.py +++ b/server/gti/gti_mcp/tools/intelligence.py @@ -68,43 +68,32 @@ async def search_iocs(query: str, ctx: Context, limit: int = 10, order_by: str = @server.tool() -async def get_hunting_ruleset(ruleset_id: str, ctx: Context) -> typing.Dict[str, typing.Any]: - """Get a Hunting Ruleset object from Google Threat Intelligence. - - A Hunting Ruleset object describes a user's hunting ruleset. It may contain multiple - Yara rules. - - The content of the Yara rules is in the `rules` attribute. - - Some important object attributes: - - creation_date: creation date as UTC timestamp. - - modification_date (int): last modification date as UTC timestamp. - - name (str): ruleset name. - - rule_names (list[str]): contains the names of all rules in the ruleset. - - number_of_rules (int): number of rules in the ruleset. - - rules (str): rule file contents. - - tags (list[str]): ruleset's custom tags. - +async def get_hunting_ruleset(query: str, ctx: Context, limit: int = 10, order_by: str = "creation_date-") -> typing.List[typing.Dict[str, typing.Any]]: + """Get hunting rulesets from Google Threat Intelligence. + Args: - ruleset_id (required): Hunting ruleset identifier. + query: Optional. A filter string to search for rulesets (e.g., "name:my_ruleset"). + limit: Optional. The maximum number of rulesets to retrieve. Defaults to 10. + order_by: Optional. The order in which to return the rulesets (e.g., "creation_date-"). Defaults to "creation_date-". Returns: - Hunting Ruleset object. + A list of Hunting Ruleset objects. """ + async with vt_client(ctx) as client: - res = await utils.fetch_object( + res = await utils.consume_vt_iterator( client, - "intelligence/hunting_rulesets", - "hunting_ruleset", - ruleset_id, - ) - return utils.sanitize_response(res) - + "/intelligence/hunting_rulesets", + params={ + "filter": query, + "order": order_by}, + limit=limit) + return utils.sanitize_response([o.to_dict() for o in res]) @server.tool() async def get_entities_related_to_a_hunting_ruleset( ruleset_id: str, relationship_name: str, ctx: Context, limit: int = 10 -) -> typing.Dict[str, typing.Any]: +) -> typing.List[typing.Dict[str, typing.Any]]: """Retrieve entities related to the the given Hunting Ruleset. The following table shows a summary of available relationships for Hunting ruleset objects. From b5db90fa43c370943563aba7500ec5d41f629093 Mon Sep 17 00:00:00 2001 From: pinksawtooth Date: Thu, 3 Jul 2025 00:09:37 +0900 Subject: [PATCH 2/2] fix get_entities_related_to_a_file --- server/gti/gti_mcp/tools/files.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/server/gti/gti_mcp/tools/files.py b/server/gti/gti_mcp/tools/files.py index da0fc35e..42441498 100644 --- a/server/gti/gti_mcp/tools/files.py +++ b/server/gti/gti_mcp/tools/files.py @@ -103,11 +103,10 @@ async def get_file_report(hash: str, ctx: Context) -> typing.Dict[str, typing.An ) return utils.sanitize_response(res) - @server.tool() async def get_entities_related_to_a_file( hash: str, relationship_name: str, descriptors_only: bool, ctx: Context, limit: int = 10, -) -> typing.Dict[str, typing.Any]: +) -> typing.Optional[typing.List[typing.Dict[str, typing.Any]]]: """Retrieve entities related to the the given file hash. The following table shows a summary of available relationships for file objects.