Skip to content

Commit 9cb90bb

Browse files
authored
chore: refactor results returning process (#36)
1 parent 65f6ce6 commit 9cb90bb

File tree

8 files changed

+205
-84
lines changed

8 files changed

+205
-84
lines changed

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ go 1.24.2
44

55
require (
66
github.com/golang-jwt/jwt/v5 v5.2.2
7-
github.com/mark3labs/mcp-go v0.23.1
7+
github.com/mark3labs/mcp-go v0.29.0
88
github.com/spf13/cobra v1.9.1
99
github.com/stretchr/testify v1.10.0
1010
github.com/zalando/go-keyring v0.2.6

go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1509,6 +1509,8 @@ github.com/mailru/easyjson v0.9.0 h1:PrnmzHw7262yW8sTBwxi1PdJA3Iw/EKBa8psRf7d9a4
15091509
github.com/mailru/easyjson v0.9.0/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU=
15101510
github.com/mark3labs/mcp-go v0.23.1 h1:RzTzZ5kJ+HxwnutKA4rll8N/pKV6Wh5dhCmiJUu5S9I=
15111511
github.com/mark3labs/mcp-go v0.23.1/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
1512+
github.com/mark3labs/mcp-go v0.29.0 h1:sH1NBcumKskhxqYzhXfGc201D7P76TVXiT0fGVhabeI=
1513+
github.com/mark3labs/mcp-go v0.29.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
15121514
github.com/masahiro331/go-disk v0.0.0-20240625071113-56c933208fee h1:cgm8mE25x5XXX2oyvJDlyJ72K+rDu/4ZCYce2worNb8=
15131515
github.com/masahiro331/go-disk v0.0.0-20240625071113-56c933208fee/go.mod h1:rojbW5tVhH1cuVYFKZS+QX+VGXK45JVsRO+jW92kkKM=
15141516
github.com/masahiro331/go-ebs-file v0.0.0-20240917043618-e6d2bea5c32e h1:nCgF1JEYIS8KNuJtIeUrmjjhktIMKWNmASZqwK2ynu0=

pkg/tools/scan/aqua.go

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package scan
33
import (
44
"context"
55
"errors"
6-
"fmt"
76
"os"
87

98
"path/filepath"
@@ -14,7 +13,7 @@ import (
1413
"github.com/mark3labs/mcp-go/mcp"
1514
)
1615

17-
func (t *ScanTools) scanWithAquaPlatform(ctx context.Context, args []string, creds creds.AquaCreds) (*mcp.CallToolResult, error) {
16+
func (t *ScanTools) scanWithAquaPlatform(ctx context.Context, args []string, creds creds.AquaCreds, scanArgs *scanArgs) (*mcp.CallToolResult, error) {
1817

1918
// add quiet to reduce the noise
2019
args = append(args, "--quiet")
@@ -87,15 +86,5 @@ func (t *ScanTools) scanWithAquaPlatform(ctx context.Context, args []string, cre
8786
return nil, err
8887
}
8988

90-
return mcp.NewToolResultResource(
91-
fmt.Sprintf(`The results can be found in the file "%s", which is found at "%s" \n
92-
Summarise the contents of the file and report it back to the user in a nicely formatted way.\n
93-
It is important that the output MUST include the ID and the severity of the issues to inform the user of the issues.
94-
`, filename, resultsFilePath),
95-
mcp.TextResourceContents{
96-
URI: resultsFilePath,
97-
MIMEType: "application/json",
98-
},
99-
), nil
100-
89+
return t.processResultsFile(resultsFilePath, scanArgs, filename)
10190
}

pkg/tools/scan/report.go

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
package scan
2+
3+
import (
4+
"encoding/json"
5+
"errors"
6+
"fmt"
7+
"os"
8+
"strings"
9+
10+
"github.com/aquasecurity/trivy/pkg/log"
11+
"github.com/aquasecurity/trivy/pkg/types"
12+
"github.com/mark3labs/mcp-go/mcp"
13+
)
14+
15+
func (t *ScanTools) processResultsFile(resultsFilePath string, scanArgs *scanArgs, filename string) (*mcp.CallToolResult, error) {
16+
logger := log.WithPrefix("scan")
17+
if scanArgs.isSBOM {
18+
// tell the LLM to present the results verbatim in code block
19+
result, err := t.processSBOMResult(resultsFilePath, logger, filename)
20+
if err != nil {
21+
logger.Error("Failed to format results", log.Err(err))
22+
return nil, fmt.Errorf("failed to format results: %w", err)
23+
}
24+
return result, nil
25+
}
26+
27+
file, err := os.Open(resultsFilePath)
28+
if err != nil {
29+
logger.Error("Failed to open scan results file", log.Err(err))
30+
return nil, errors.New("failed to open scan results file")
31+
}
32+
defer func() {
33+
if err := file.Close(); err != nil {
34+
logger.Error("Failed to close scan results file", log.Err(err))
35+
}
36+
if err := os.Remove(resultsFilePath); err != nil {
37+
logger.Error("Failed to remove scan results file", log.Err(err))
38+
}
39+
}()
40+
41+
var rep types.Report
42+
if err := json.NewDecoder(file).Decode(&rep); err != nil {
43+
logger.Error("Failed to decode scan results file", log.Err(err))
44+
return nil, errors.New("failed to decode scan results file")
45+
}
46+
47+
// check the size of the file, if its larger than 1MB, we don't want to embed it in the response
48+
// instead we want to provide a link to the file
49+
fileInfo, err := file.Stat()
50+
if err != nil {
51+
logger.Error("Failed to get scan results file info", log.Err(err))
52+
return nil, errors.New("failed to get scan results file info")
53+
}
54+
55+
var resultString strings.Builder
56+
57+
if fileInfo.Size() > 1024*1024 {
58+
resultString, err = processResultSummary(rep, logger)
59+
if err != nil {
60+
logger.Error("Failed to format results", log.Err(err))
61+
return nil, fmt.Errorf("failed to format results: %w", err)
62+
}
63+
} else {
64+
resultString, err = processResult(rep, logger)
65+
if err != nil {
66+
logger.Error("Failed to format results", log.Err(err))
67+
return nil, fmt.Errorf("failed to format results: %w", err)
68+
}
69+
}
70+
return mcp.NewToolResultText(resultString.String()), nil
71+
}
72+
73+
func processResultSummary(rep types.Report, logger *log.Logger) (strings.Builder, error) {
74+
// process the results into a text format
75+
sb := strings.Builder{}
76+
77+
logger.Debug("Scan results file is larger than 1MB, building a summary of the results")
78+
countsMap := make(map[string]int)
79+
80+
for _, result := range rep.Results {
81+
countsMap["Vulnerabilities"] += len(result.Vulnerabilities)
82+
countsMap["Misconfigurations"] += len(result.Misconfigurations)
83+
countsMap["Licenses"] += len(result.Licenses)
84+
countsMap["Secrets"] += len(result.Secrets)
85+
}
86+
87+
sb.WriteString("## Scan Results Summary\n")
88+
sb.WriteString(fmt.Sprintf(" - Vulnerabilities: %d\n", countsMap["Vulnerabilities"]))
89+
sb.WriteString(fmt.Sprintf(" - Misconfigurations: %d\n", countsMap["Misconfigurations"]))
90+
sb.WriteString(fmt.Sprintf(" - Licenses: %d\n", countsMap["Licenses"]))
91+
sb.WriteString(fmt.Sprintf(" - Secrets: %d\n", countsMap["Secrets"]))
92+
sb.WriteString("\n\n")
93+
94+
return sb, nil
95+
96+
}
97+
98+
func processResult(rep types.Report, logger *log.Logger) (strings.Builder, error) {
99+
100+
// process the results into a text format
101+
sb := strings.Builder{}
102+
sb.WriteString("The scan results are below, it is important that you present the results with the severity and the ID/Name of the vulnerability/misconfiguration/license/secret. \n")
103+
104+
for _, result := range rep.Results {
105+
if len(result.Vulnerabilities) > 0 || len(result.Misconfigurations) > 0 || len(result.Licenses) > 0 || len(result.Secrets) > 0 {
106+
sb.WriteString(fmt.Sprintf("## %s\n", result.Target))
107+
for _, vuln := range result.Vulnerabilities {
108+
sb.WriteString(fmt.Sprintf("### %s\n", vuln.VulnerabilityID))
109+
sb.WriteString(fmt.Sprintf(" - Severity: %s\n", vuln.Severity))
110+
sb.WriteString(fmt.Sprintf(" - Package: %s\n", vuln.PkgName))
111+
sb.WriteString(fmt.Sprintf(" - Installed Version: %s\n", vuln.InstalledVersion))
112+
sb.WriteString(fmt.Sprintf(" - Fixed Version: %s\n", vuln.FixedVersion))
113+
sb.WriteString(fmt.Sprintf(" - Primary URL: %s\n", vuln.PrimaryURL))
114+
sb.WriteString(fmt.Sprintf(" - Data Source: %s\n", vuln.DataSource))
115+
}
116+
for _, misconf := range result.Misconfigurations {
117+
sb.WriteString(fmt.Sprintf("### %s\n", misconf.ID))
118+
sb.WriteString(fmt.Sprintf(" - Severity: %s\n", misconf.Severity))
119+
sb.WriteString(fmt.Sprintf(" - Title: %s\n", misconf.Title))
120+
sb.WriteString(fmt.Sprintf(" - Description: %s\n", misconf.Description))
121+
sb.WriteString(fmt.Sprintf(" - Resolution: %s\n", misconf.Resolution))
122+
sb.WriteString(fmt.Sprintf(" - Primary URL: %s\n", misconf.PrimaryURL))
123+
sb.WriteString(fmt.Sprintf(" - Message: %s\n", misconf.Message))
124+
125+
}
126+
for _, license := range result.Licenses {
127+
sb.WriteString(fmt.Sprintf("### %s\n", license.Name))
128+
sb.WriteString(fmt.Sprintf(" - Severity: %s\n", license.Severity))
129+
sb.WriteString(fmt.Sprintf(" - Package: %s\n", license.PkgName))
130+
sb.WriteString(fmt.Sprintf(" - License: %s\n", license.Text))
131+
sb.WriteString(fmt.Sprintf(" - Category: %s\n", license.Category))
132+
}
133+
for _, secret := range result.Secrets {
134+
sb.WriteString(fmt.Sprintf("### %s\n", secret.RuleID))
135+
sb.WriteString(fmt.Sprintf(" - Severity: %s\n", secret.Severity))
136+
sb.WriteString(fmt.Sprintf(" - Title: %s\n", secret.Title))
137+
sb.WriteString(fmt.Sprintf(" - Code: %v\n", secret.Code))
138+
sb.WriteString(fmt.Sprintf(" - Match: %s\n", secret.Match))
139+
sb.WriteString(fmt.Sprintf(" - Category: %s\n", secret.Category))
140+
}
141+
}
142+
143+
}
144+
return sb, nil
145+
}
146+
147+
// processSBOMResult processes the SBOM result and returns a tool result
148+
// we don't clean up the results file here because we want to keep it to be available for the LLM to provide a link
149+
// when the MCP server is closed, the trivy mcp cache should be cleaned up
150+
func (*ScanTools) processSBOMResult(resultsFilePath string, logger *log.Logger, filename string) (*mcp.CallToolResult, error) {
151+
log.Debug("Scan results file", log.String("file", resultsFilePath))
152+
153+
// Check that the file exists
154+
if _, err := os.Stat(resultsFilePath); err != nil {
155+
logger.Error("Failed to access scan results file", log.Err(err))
156+
return nil, errors.New("failed to access scan results file")
157+
}
158+
159+
// Create a markdown link for the SBOM file instead of using resource object
160+
// This provides better compatibility with VSCode's MCP client
161+
markdownLink := fmt.Sprintf("Only show the user the following markdown and NOTHING else \n [SBOM Results](file://%s)", resultsFilePath)
162+
163+
// For large files, we return a link that client can click instead of embedding the content
164+
// This avoids issues with resource handling in VSCode's MCP client
165+
return mcp.NewToolResultText(markdownLink), nil
166+
}

pkg/tools/scan/scan.go

Lines changed: 17 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import (
1616
"github.com/aquasecurity/trivy/pkg/log"
1717
"github.com/mark3labs/mcp-go/mcp"
1818

19-
_ "modernc.org/sqlite" // sqlite driver for RPM DB and Java DB
19+
_ "modernc.org/sqlite" // sqlite driver for RPM DB and Java DB, this is needed for the scan
2020
)
2121

2222
type ScanTools struct {
@@ -38,44 +38,47 @@ func NewScanTools(opts flag.Options, trivyTempDir string) *ScanTools {
3838
var (
3939
ScanFilesystemTool = mcp.NewTool(
4040
"scan_filesystem",
41-
mcp.WithDescription(`Scan a local filesystem project for vulnerabilities, misconfigurations, licenses, and secrets issue using Trivy. \n
42-
When the result is an embedded resource (like an SBOM), you MUST format the response as a clickable markdown link with the text set to "SBOM Results" and with the URI as the target of the link. \n
43-
Do not include any other text or explanation with the link.".`),
41+
mcp.WithDescription(`Scan a local filesystem project for vulnerabilities, misconfigurations, licenses, and secrets issue using Trivy.
42+
Follow the instructions that are given in the response.`),
4443
targetString,
4544
scanTypeArray,
4645
severityArray,
4746
outputFormatString,
4847
fixedOnlyBool,
4948
targetTypeString("filesystem"),
5049
mcp.WithToolAnnotation(mcp.ToolAnnotation{
51-
Title: "Scan filesystem and local projects with Trivy",
50+
Title: "Scan local filesystem with Trivy",
5251
}),
5352
)
5453

5554
ScanImageTool = mcp.NewTool(
5655
"scan_image",
57-
mcp.WithDescription(`Scan a container image for vulnerabilities, misconfigurations, licenses, and secrets issue using Trivy \n
58-
When the result is an embedded resource (like an SBOM), you MUST format the response as a clickable markdown link with the text set to the filename and the URI as the target of the link. \n
59-
Do not include any other text or explanation with the link.".`),
56+
mcp.WithDescription(`Scan a container image for vulnerabilities, misconfigurations, licenses, and secrets issue using Trivy
57+
Follow the instructions that are given in the response.`),
6058
targetString,
6159
scanTypeArray,
6260
severityArray,
6361
outputFormatString,
6462
fixedOnlyBool,
6563
targetTypeString("image"),
64+
mcp.WithToolAnnotation(mcp.ToolAnnotation{
65+
Title: "Scan a container image with Trivy",
66+
}),
6667
)
6768

6869
ScanRepositoryTool = mcp.NewTool(
6970
"scan_repository",
70-
mcp.WithDescription(`Scan a remote git repository for vulnerabilities, misconfigurations, licenses, and secrets issue using Trivy \n
71-
When the result is an embedded resource (like an SBOM), you MUST format the response as a clickable markdown link with the text set to "SBOM Results" and with the URI as the target of the link. \n
72-
Do not include any other text or explanation with the link.".`),
71+
mcp.WithDescription(`Scan a remote git repository for vulnerabilities, misconfigurations, licenses, and secrets issue using Trivy.
72+
Follow the instructions that are given in the response.`),
7373
targetString,
7474
scanTypeArray,
7575
severityArray,
7676
outputFormatString,
7777
fixedOnlyBool,
7878
targetTypeString("repository"),
79+
mcp.WithToolAnnotation(mcp.ToolAnnotation{
80+
Title: "Scan a remote git repository with Trivy",
81+
}),
7982
)
8083
)
8184

@@ -101,7 +104,7 @@ func (t *ScanTools) ScanWithTrivyHandler(ctx context.Context, request mcp.CallTo
101104
args = append(args, "--skip-update")
102105
}
103106

104-
// json output doesn't include the target in the output
107+
// json output doesn't include the packages in the output
105108
if scanArgs.outputFormat == "json" && slices.Contains(scanArgs.scanType, "vuln") {
106109
args = append(args, "--list-all-pkgs")
107110
}
@@ -113,7 +116,7 @@ func (t *ScanTools) ScanWithTrivyHandler(ctx context.Context, request mcp.CallTo
113116
return nil, fmt.Errorf("failed to load credentials which suggests the haven't been saved using `trivy mcp auth`: %v", err)
114117
}
115118
args = append(args, scanArgs.target)
116-
return t.scanWithAquaPlatform(ctx, args, *aquaCreds)
119+
return t.scanWithAquaPlatform(ctx, args, *aquaCreds, scanArgs)
117120
}
118121

119122
logger := log.WithPrefix(scanArgs.targetType)
@@ -150,47 +153,8 @@ func (t *ScanTools) ScanWithTrivyHandler(ctx context.Context, request mcp.CallTo
150153

151154
logger.Info("Scan completed successfully")
152155

153-
if scanArgs.isSBOM {
154-
// tell the LLM to present the results verbatim in code block
155-
result, err := t.processSBOMResult(resultsFilePath, logger, filename)
156-
if err != nil {
157-
logger.Error("Failed to format results", log.Err(err))
158-
return nil, fmt.Errorf("failed to format results: %w", err)
159-
}
160-
return result, nil
161-
}
162-
163-
return mcp.NewToolResultResource(
164-
fmt.Sprintf(`The results can be found in the file "%s", which is found at "%s" \n
165-
Summarise the contents of the file and report it back to the user in a nicely formatted way.\n
166-
It is important that the output MUST include the ID and the severity of the issues to inform the user of the issues.
167-
`, filename, resultsFilePath),
168-
mcp.TextResourceContents{
169-
URI: resultsFilePath,
170-
MIMEType: "application/json",
171-
},
172-
), nil
173-
}
174-
175-
// processSBOMResult processes the SBOM result and returns a tool result
176-
// we don't clean up the results file here because we want to keep it to be available for the LLM to provide a link
177-
// when the MCP server is closed, the trivy mcp cache should be cleaned up
178-
func (*ScanTools) processSBOMResult(resultsFilePath string, logger *log.Logger, filename string) (*mcp.CallToolResult, error) {
179-
log.Debug("Scan results file", log.String("file", resultsFilePath))
180-
181-
content, err := os.ReadFile(resultsFilePath)
182-
if err != nil {
183-
logger.Error("Failed to read scan results file", log.Err(err))
184-
return nil, errors.New("failed to read scan results file")
185-
}
156+
return t.processResultsFile(resultsFilePath, scanArgs, filename)
186157

187-
return mcp.NewToolResultResource(
188-
fmt.Sprintf("The embedded resource is a human readable SBOM format. \nThe filename is %s and the URI is file://%s. \n", filename, resultsFilePath),
189-
mcp.TextResourceContents{
190-
URI: resultsFilePath,
191-
Text: string(content),
192-
},
193-
), nil
194158
}
195159

196160
func getFilename(targetType, format string) string {

0 commit comments

Comments
 (0)