From 6e3447f89ec92dae3eec7fa9d2084a010deb4a09 Mon Sep 17 00:00:00 2001 From: Spyros Date: Thu, 28 Apr 2022 20:53:34 +0100 Subject: [PATCH] Rewrite stdout json to go (#129) * make stdout_json consumer in go * cleanup * fix docker deps --- consumers/stdout_json/BUILD | 23 +--- consumers/stdout_json/Dockerfile | 6 +- consumers/stdout_json/main.go | 159 ++++++++++++++++++++++ consumers/stdout_json/stdout_json.py | 100 -------------- consumers/stdout_json/stdout_json_test.py | 139 ------------------- 5 files changed, 168 insertions(+), 259 deletions(-) create mode 100644 consumers/stdout_json/main.go delete mode 100644 consumers/stdout_json/stdout_json.py delete mode 100644 consumers/stdout_json/stdout_json_test.py diff --git a/consumers/stdout_json/BUILD b/consumers/stdout_json/BUILD index d655be8..2a9be2b 100644 --- a/consumers/stdout_json/BUILD +++ b/consumers/stdout_json/BUILD @@ -1,33 +1,22 @@ subinclude("//third_party/defs:docker") -python_binary( +go_binary( name = "stdout_json", - main = "stdout_json.py", - deps = [ - "//api/proto:v1", - "//consumers:consumers_base_python", - "//third_party/python:protobuf", - ], -) - -python_test( - name = "stdout_json_test", srcs = [ - "stdout_json_test.py", + "main.go", ], deps = [ - ":stdout_json", "//api/proto:v1", - "//consumers:consumers_base_python", - "//third_party/python:protobuf", + "//consumers", + "//third_party/go:protobuf", ], ) docker_image( - name = "dracon-consumer-stdout", + name = "image", srcs = [ ":stdout_json", ], - base_image = "//build/docker:dracon-base-python", + base_image = "//build/docker:dracon-base-go", image = "dracon-consumer-stdout-json", ) diff --git a/consumers/stdout_json/Dockerfile b/consumers/stdout_json/Dockerfile index 1bbaeb9..2720291 100644 --- a/consumers/stdout_json/Dockerfile +++ b/consumers/stdout_json/Dockerfile @@ -1,5 +1,5 @@ -FROM //build/docker:dracon-base-python +FROM //build/docker:dracon-base-go -COPY /stdout_json.pex /consume.pex +COPY stdout_json /consume -ENTRYPOINT ["/consume.pex"] +ENTRYPOINT ["/consume"] diff --git a/consumers/stdout_json/main.go b/consumers/stdout_json/main.go new file mode 100644 index 0000000..d3951bf --- /dev/null +++ b/consumers/stdout_json/main.go @@ -0,0 +1,159 @@ +package main + +import ( + "encoding/json" + "fmt" + "log" + "time" + + "github.com/golang/protobuf/ptypes" + v1 "github.com/thought-machine/dracon/api/proto/v1" + "github.com/thought-machine/dracon/consumers" +) + +func parseFlags() error { + if err := consumers.ParseFlags(); err != nil { + return err + } + return nil +} + +func main() { + if err := consumers.ParseFlags(); err != nil { + log.Fatal(err) + } + if consumers.Raw { + responses, err := consumers.LoadToolResponse() + if err != nil { + log.Fatal("could not load raw results, file malformed: ", err) + } + for _, res := range responses { + scanStartTime, _ := ptypes.Timestamp(res.GetScanInfo().GetScanStartTime()) + for _, iss := range res.GetIssues() { + b, err := getRawIssue(scanStartTime, res, iss) + if err != nil { + log.Fatal("Could not parse raw issue", err) + } + fmt.Printf("%s", string(b)) + } + } + } else { + responses, err := consumers.LoadEnrichedToolResponse() + if err != nil { + log.Fatal("could not load enriched results, file malformed: ", err) + } + for _, res := range responses { + scanStartTime, _ := ptypes.Timestamp(res.GetOriginalResults().GetScanInfo().GetScanStartTime()) + for _, iss := range res.GetIssues() { + b, err := getEnrichedIssue(scanStartTime, res, iss) + if err != nil { + log.Fatal("Could not parse enriched issue", err) + } + fmt.Printf("%s", string(b)) + } + } + } +} + +func getRawIssue(scanStartTime time.Time, res *v1.LaunchToolResponse, iss *v1.Issue) ([]byte, error) { + jBytes, err := json.Marshal(&draconDocument{ + ScanStartTime: scanStartTime, + ScanID: res.GetScanInfo().GetScanUuid(), + ToolName: res.GetToolName(), + Source: iss.GetSource(), + Title: iss.GetTitle(), + Target: iss.GetTarget(), + Type: iss.GetType(), + Severity: iss.GetSeverity(), + CVSS: iss.GetCvss(), + Confidence: iss.GetConfidence(), + Description: iss.GetDescription(), + FirstFound: scanStartTime, + Count: 1, + FalsePositive: false, + CVE: iss.GetCve(), + }) + if err != nil { + return []byte{}, err + } + return jBytes, nil +} +func severtiyToText(severity v1.Severity) string { + switch severity { + case v1.Severity_SEVERITY_INFO: + return "Info" + case v1.Severity_SEVERITY_LOW: + return "Low" + case v1.Severity_SEVERITY_MEDIUM: + return "Medium" + case v1.Severity_SEVERITY_HIGH: + return "High" + case v1.Severity_SEVERITY_CRITICAL: + return "Critical" + default: + return "N/A" + } +} +func confidenceToText(confidence v1.Confidence) string { + switch confidence { + case v1.Confidence_CONFIDENCE_INFO: + return "Info" + case v1.Confidence_CONFIDENCE_LOW: + return "Low" + case v1.Confidence_CONFIDENCE_MEDIUM: + return "Medium" + case v1.Confidence_CONFIDENCE_HIGH: + return "High" + case v1.Confidence_CONFIDENCE_CRITICAL: + return "Critical" + default: + return "N/A" + } + +} +func getEnrichedIssue(scanStartTime time.Time, res *v1.EnrichedLaunchToolResponse, iss *v1.EnrichedIssue) ([]byte, error) { + firstSeenTime, _ := ptypes.Timestamp(iss.GetFirstSeen()) + jBytes, err := json.Marshal(&draconDocument{ + ScanStartTime: scanStartTime, + ScanID: res.GetOriginalResults().GetScanInfo().GetScanUuid(), + ToolName: res.GetOriginalResults().GetToolName(), + Source: iss.GetRawIssue().GetSource(), + Title: iss.GetRawIssue().GetTitle(), + Target: iss.GetRawIssue().GetTarget(), + Type: iss.GetRawIssue().GetType(), + Severity: iss.GetRawIssue().GetSeverity(), + CVSS: iss.GetRawIssue().GetCvss(), + Confidence: iss.GetRawIssue().GetConfidence(), + Description: iss.GetRawIssue().GetDescription(), + FirstFound: firstSeenTime, + Count: iss.GetCount(), + FalsePositive: iss.GetFalsePositive(), + SeverityText: severtiyToText(iss.GetRawIssue().GetSeverity()), + ConfidenceText: confidenceToText(iss.GetRawIssue().GetConfidence()), + CVE: iss.GetRawIssue().GetCve(), + }) + if err != nil { + return []byte{}, err + } + return jBytes, nil +} + +type draconDocument struct { + ScanStartTime time.Time `json:"scan_start_time"` + ScanID string `json:"scan_id"` + ToolName string `json:"tool_name"` + Source string `json:"source"` + Target string `json:"target"` + Type string `json:"type"` + Title string `json:"title"` + Severity v1.Severity `json:"severity"` + SeverityText string `json:"severity_text"` + CVSS float64 `json:"cvss"` + Confidence v1.Confidence `json:"confidence"` + ConfidenceText string `json:"confidence_text"` + Description string `json:"description"` + FirstFound time.Time `json:"first_found"` + Count uint64 `json:"count"` + FalsePositive bool `json:"false_positive"` + CVE string `json:"cve"` +} diff --git a/consumers/stdout_json/stdout_json.py b/consumers/stdout_json/stdout_json.py deleted file mode 100644 index a8c6e6f..0000000 --- a/consumers/stdout_json/stdout_json.py +++ /dev/null @@ -1,100 +0,0 @@ -import argparse -import sys -import json -from api.proto import engine_pb2 -from consumers.consumer import Consumer -import logging - - -logger = logging.getLogger(__name__) - -class StdoutJsonConsumer(Consumer): - - def __init__(self, config: dict): - logger.info("Starting Consumer") - self.pvc_location = config.pvc_location - logger.info("Reading from %s" % self.pvc_location) - - if (self.pvc_location is None): - raise AttributeError("PVC claim location is missing") - - def load_results(self) -> (list, bool): - try: - return self._load_enriched_results(), False - except SyntaxError: - return self._load_plain_results(), True - - def _load_plain_results(self): - scan_results = engine_pb2.LaunchToolResponse() - return self.load_files(scan_results, self.pvc_location) - - def _load_enriched_results(self): - """Load a set of LaunchToolResponse protobufs into a list for processing""" - return super().load_results() - - def print_data(self, data): - print(json.dumps(data)) - - def send_results(self, collected_results: list, raw: bool): - """ - Take a list of LaunchToolResponse protobufs and print them - - :param collected_results: list of LaunchToolResponse protobufs - """ - for sc in collected_results: - for iss in sc.issues: - if raw: - scan = sc - issue = iss - first_found = scan.scan_info.scan_start_time.ToJsonString() - count = 1 - false_positive = False - else: - issue = iss.raw_issue - first_found = iss.first_seen.ToJsonString() - count = iss.count - false_positive = iss.false_positive - scan = sc.original_results - data = { - 'scan_start_time': scan.scan_info.scan_start_time.ToJsonString(), - 'scan_id': scan.scan_info.scan_uuid, - 'tool_name': scan.tool_name, - 'target': issue.target, - 'type': issue.type, - 'title': issue.title, - 'severity': issue.severity, - 'cvss': issue.cvss, - 'confidence': issue.confidence, - 'description': issue.description, - 'first_found': first_found, - 'count': count, - 'false_positive': false_positive, - 'cve': issue.cve, - } - self.print_data(data) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - '--pvc_location', help='The location of the scan results') - parser.add_argument( - '--raw', help='if it should process raw or enriched results', action="store_true") - args = parser.parse_args() - ec = StdoutJsonConsumer(args) - try: - logger.info('Loading results from %s' % str(ec.pvc_location)) - collected_results, raw = ec.load_results() - logger.info("gathered %s results", len(collected_results)) - logger.info("Reading raw: %s ", raw) - except SyntaxError as e: - logger.error('Unable to load results from %s: ', str(e)) - sys.exit(-1) - - ec.send_results(collected_results, args.raw) - logger.info('Done!') - - -if __name__ == '__main__': - logger.info("Consumer Stdout JSON running") - main() diff --git a/consumers/stdout_json/stdout_json_test.py b/consumers/stdout_json/stdout_json_test.py deleted file mode 100644 index f43e1c8..0000000 --- a/consumers/stdout_json/stdout_json_test.py +++ /dev/null @@ -1,139 +0,0 @@ -import unittest -import tempfile -import shutil -import logging - -from third_party.python.google.protobuf.timestamp_pb2 import Timestamp -from unittest import mock -from consumers.stdout_json import stdout_json -from api.proto import engine_pb2 -from api.proto import issue_pb2 - -logger = logging.getLogger(__name__) - - -class Config: - pvc_location = '/tmp/' - -class TestJSONConsumer(unittest.TestCase): - - def setUp(self): - self.config = Config() - scan_start_time = Timestamp() - scan_start_time.FromJsonString("1991-01-01T00:00:00Z") - scan_info = engine_pb2.ScanInfo( - scan_start_time=scan_start_time, - scan_uuid='dd1794f2-544d-456b-a45a-a2bec53633b1' - ) - scan_results = engine_pb2.LaunchToolResponse( - scan_info=scan_info - ) - scan_results.tool_name = 'unit_tests' - - issue = issue_pb2.Issue() - issue.target = 'target.py:0' - issue.type = "test" - issue.title = "test title" - issue.cvss = 2.0 - issue.description = "test.description" - issue.severity = issue_pb2.Severity.SEVERITY_LOW - issue.confidence = issue_pb2.Confidence.CONFIDENCE_LOW - issue.cve = "CVE-0000-99999" - - scan_results.issues.extend([issue]) - first_seen = Timestamp() - first_seen.FromJsonString("1992-02-02T00:00:00Z") - enriched_issue = issue_pb2.EnrichedIssue(first_seen=first_seen) - enriched_issue.raw_issue.CopyFrom(issue) - enriched_issue.count = 2 - enriched_issue.false_positive = True - - enriched_scan_results = engine_pb2.EnrichedLaunchToolResponse( - original_results=scan_results, - ) - enriched_scan_results.issues.extend([enriched_issue]) - - self.enriched_dtemp = tempfile.mkdtemp( - prefix="enriched_", dir=self.config.pvc_location) - self.enriched, _ = tempfile.mkstemp( - prefix="enriched_", dir=self.enriched_dtemp, suffix=".pb") - - self.raw_dtemp = tempfile.mkdtemp( - prefix="raw_", dir=self.config.pvc_location) - self.raw, _ = tempfile.mkstemp( - prefix="raw_", dir=self.raw_dtemp, suffix=".pb") - - f = open(self.enriched, "wb") - scan_proto_string = enriched_scan_results.SerializeToString() - f.write(scan_proto_string) - f.close() - - f = open(self.raw, "wb") - scan_proto_string = scan_results.SerializeToString() - f.write(scan_proto_string) - f.close() - - def tearDown(self): - shutil.rmtree(self.raw_dtemp) - shutil.rmtree(self.enriched_dtemp) - - def _create_consumer(self): - return stdout_json.StdoutJsonConsumer(self.config) - - def test_load_results(self): - self.config.pvc_location = self.enriched_dtemp - consumer = stdout_json.StdoutJsonConsumer(self.config) - _, raw = consumer.load_results() - self.assertFalse(raw) - - self.config.pvc_location = self.raw_dtemp - consumer = stdout_json.StdoutJsonConsumer(self.config) - _, raw = consumer.load_results() - self.assertTrue(raw) - - @mock.patch( - "consumers.stdout_json.stdout_json" - ".StdoutJsonConsumer" - ".print_data") - def test_send_results(self, mocked_print_data): - mocked_print_data.return_value = "" - - data = { - 'scan_start_time': "1991-01-01T00:00:00Z", - 'scan_id': 'dd1794f2-544d-456b-a45a-a2bec53633b1', - 'tool_name': 'unit_tests', - 'target': 'target.py:0', - 'type': "test", - 'title': "test title", - 'severity': 1, - 'cvss': 2.0, - 'confidence': 1, - 'description': "test.description", - 'first_found': '', - 'count': 1, - 'false_positive': '', - 'cve': 'CVE-0000-99999', - } - - self.config.pvc_location = self.raw_dtemp - consumer = stdout_json.StdoutJsonConsumer(self.config) - results, raw = consumer.load_results() - - consumer.send_results(results, raw) - data['first_found'] = data['scan_start_time'] - data['false_positive'] = False - mocked_print_data.assert_called_with(data) - - self.config.pvc_location = self.enriched_dtemp - consumer = stdout_json.StdoutJsonConsumer(self.config) - results, raw = consumer.load_results() - - consumer.send_results(results, raw) - data['first_found'] = "1992-02-02T00:00:00Z" - data['count'] = 2 - data['false_positive'] = True - mocked_print_data.assert_called_with(data) - - -if __name__ == '__main__': - unittest.main()