diff --git a/deepfence_server/handler/export_reports.go b/deepfence_server/handler/export_reports.go index cbd2b9e460..ba26d7a198 100644 --- a/deepfence_server/handler/export_reports.go +++ b/deepfence_server/handler/export_reports.go @@ -260,10 +260,11 @@ func (h *Handler) GenerateReport(w http.ResponseWriter, r *http.Request) { // report task params report_id := uuid.New().String() params := utils.ReportParams{ - ReportID: report_id, - ReportType: req.ReportType, - Duration: req.Duration, - Filters: req.Filters, + ReportID: report_id, + ReportType: req.ReportType, + Duration: req.Duration, + Filters: req.Filters, + CustomFields: req.CustomFields, } worker, err := directory.Worker(r.Context()) @@ -294,14 +295,15 @@ func (h *Handler) GenerateReport(w http.ResponseWriter, r *http.Request) { defer tx.Close() query := ` - CREATE (n:Report{created_at:TIMESTAMP(), type:$type, report_id:$uid, status:$status, filters:$filters, duration:$duration}) + CREATE (n:Report{created_at:TIMESTAMP(), type:$type, report_id:$uid, status:$status, filters:$filters, duration:$duration, custom_fields:$custom_fields}) RETURN n` vars := map[string]interface{}{ - "type": req.ReportType, - "uid": report_id, - "status": utils.SCAN_STATUS_STARTING, - "filters": req.Filters.String(), - "duration": req.Duration, + "type": req.ReportType, + "uid": report_id, + "status": utils.SCAN_STATUS_STARTING, + "filters": req.Filters.String(), + "duration": req.Duration, + "custom_fields": req.CustomFields, } _, err = tx.Run(query, vars) diff --git a/deepfence_server/model/reports.go b/deepfence_server/model/reports.go index f10cefdf21..bdd37880f5 100644 --- a/deepfence_server/model/reports.go +++ b/deepfence_server/model/reports.go @@ -5,9 +5,10 @@ import ( ) type GenerateReportReq struct { - ReportType string `json:"report_type" validate:"required" required:"true" enum:"pdf,xlsx"` - Duration int `json:"duration" enum:"0,1,7,30,60,90,180"` - Filters utils.ReportFilters `json:"filters"` + ReportType string `json:"report_type" validate:"required" required:"true" enum:"pdf,xlsx"` + Duration int `json:"duration" enum:"0,1,7,30,60,90,180"` + Filters utils.ReportFilters `json:"filters"` + CustomFields []string `json:"custom_fields"` } type GenerateReportResp struct { diff --git a/deepfence_utils/utils/structs.go b/deepfence_utils/utils/structs.go index b9eb7038a7..d769db0e8c 100644 --- a/deepfence_utils/utils/structs.go +++ b/deepfence_utils/utils/structs.go @@ -71,10 +71,11 @@ type MalwareScanParameters struct { } type ReportParams struct { - ReportID string `json:"report_id"` - ReportType string `json:"report_type"` - Duration int `json:"duration"` - Filters ReportFilters `json:"filters"` + ReportID string `json:"report_id"` + ReportType string `json:"report_type"` + Duration int `json:"duration"` + Filters ReportFilters `json:"filters"` + CustomFields []string `json:"custom_fields"` } type ReportFilters struct { diff --git a/deepfence_worker/tasks/reports/xlsx.go b/deepfence_worker/tasks/reports/xlsx.go index 6c4f40ce84..0a36cdecc6 100644 --- a/deepfence_worker/tasks/reports/xlsx.go +++ b/deepfence_worker/tasks/reports/xlsx.go @@ -3,6 +3,7 @@ package reports import ( "context" "os" + "reflect" "github.com/deepfence/ThreatMapper/deepfence_utils/log" "github.com/deepfence/ThreatMapper/deepfence_utils/utils" @@ -72,6 +73,20 @@ var ( } ) +func customFieldsToMap(customFields []string) map[string]string { + result := make(map[string]string) + + for i, field := range customFields { + // Convert index to corresponding column letter + colLetter := string(rune('A' + i)) + + // Create map entry with the column letter as key and field as value + result[colLetter+"1"] = field + } + + return result +} + func generateXLSX(ctx context.Context, params utils.ReportParams) (string, error) { var ( @@ -126,6 +141,45 @@ func xlsxSetHeader(xlsx *excelize.File, sheet string, headers map[string]string) } } +func getMatchingValues(v interface{}, scanInfo interface{}, customFields []string) ([]interface{}, error) { + var matchingValues []interface{} + + vValue := reflect.ValueOf(v) + scanInfoValue := reflect.ValueOf(scanInfo) + + for _, field := range customFields { + // Get field values using reflection and JSON tags + vFieldValue := getFieldByJSONTag(vValue, field) + scanInfoFieldValue := getFieldByJSONTag(scanInfoValue, field) + + // Check which field has valid value + if vFieldValue.IsValid() && !vFieldValue.IsZero() { + matchingValues = append(matchingValues, vFieldValue.Interface()) + continue + } + + if scanInfoFieldValue.IsValid() && !scanInfoFieldValue.IsZero() { + matchingValues = append(matchingValues, scanInfoFieldValue.Interface()) + continue + } + } + + return matchingValues, nil +} + +// getFieldByJSONTag retrieves a field value by its JSON tag +func getFieldByJSONTag(value reflect.Value, jsonTag string) reflect.Value { + for i := 0; i < value.NumField(); i++ { + field := value.Type().Field(i) + tag := field.Tag.Get("json") + + if tag == jsonTag { + return value.Field(i) + } + } + return reflect.Value{} +} + func vulnerabilityXLSX(ctx context.Context, params utils.ReportParams) (string, error) { data, err := getVulnerabilityData(ctx, params) if err != nil { @@ -140,6 +194,12 @@ func vulnerabilityXLSX(ctx context.Context, params utils.ReportParams) (string, } }() + if len(params.CustomFields) > 0 { + // make vulnerabilityHeader according to custom fields + vulnerabilityHeader = customFieldsToMap(params.CustomFields) + } + log.Info().Msgf("vulnerabilityHeader: %+v", vulnerabilityHeader) + xlsxSetHeader(xlsx, "Sheet1", vulnerabilityHeader) offset := 0 @@ -149,24 +209,32 @@ func vulnerabilityXLSX(ctx context.Context, params utils.ReportParams) (string, if err != nil { log.Error().Err(err).Msg("error generating cell name") } - value := []interface{}{ - nodeScanData.ScanInfo.UpdatedAt, - v.Cve_attack_vector, - v.Cve_caused_by_package, - nodeScanData.ScanInfo.NodeName, - nodeScanData.ScanInfo.ScanID, - nodeScanData.ScanInfo.NodeID, - v.Cve_cvss_score, - v.Cve_description, - v.Cve_fixed_in, - v.Cve_id, - v.Cve_link, - v.Cve_severity, - v.Cve_overall_score, - v.Cve_type, - nodeScanData.ScanInfo.HostName, - nodeScanData.ScanInfo.CloudAccountID, - v.Masked, + var value []interface{} + if len(params.CustomFields) > 0 { + value, err = getMatchingValues(v, nodeScanData.ScanInfo, params.CustomFields) + if err != nil { + log.Error().Err(err).Msg("error getting matching values") + } + } else { + value = []interface{}{ + nodeScanData.ScanInfo.UpdatedAt, + v.Cve_attack_vector, + v.Cve_caused_by_package, + nodeScanData.ScanInfo.NodeName, + nodeScanData.ScanInfo.ScanID, + nodeScanData.ScanInfo.NodeID, + v.Cve_cvss_score, + v.Cve_description, + v.Cve_fixed_in, + v.Cve_id, + v.Cve_link, + v.Cve_severity, + v.Cve_overall_score, + v.Cve_type, + nodeScanData.ScanInfo.HostName, + nodeScanData.ScanInfo.CloudAccountID, + v.Masked, + } } err = xlsx.SetSheetRow("Sheet1", cellName, &value) if err != nil {