From c43bda20fc2f8189bfaa39322ba6e1c644198a56 Mon Sep 17 00:00:00 2001 From: sg Date: Thu, 24 Oct 2024 15:27:58 +0100 Subject: [PATCH] fixup --- components/consumers/pdf/main_test.go | 62 +++++++++++++++++++++++++ pkg/playwright/mock/wrapper.go | 19 ++++++++ pkg/playwright/wrapper.go | 58 ++++++++++++++++++++++++ pkg/s3/mock/wrapper.go | 16 +++++++ pkg/s3/wrapper.go | 65 +++++++++++++++++++++++++++ 5 files changed, 220 insertions(+) create mode 100644 components/consumers/pdf/main_test.go create mode 100644 pkg/playwright/mock/wrapper.go create mode 100644 pkg/playwright/wrapper.go create mode 100644 pkg/s3/mock/wrapper.go create mode 100644 pkg/s3/wrapper.go diff --git a/components/consumers/pdf/main_test.go b/components/consumers/pdf/main_test.go new file mode 100644 index 000000000..d76744045 --- /dev/null +++ b/components/consumers/pdf/main_test.go @@ -0,0 +1,62 @@ +// Package main of the pdf consumer implements a simple consumer for +// applying a go-template to a smithy scan, converting the result to pdf and then +// uploading the result to the S3 bucket passed as an argument +// the consumer expects the environment variables +// AWS_ACCESS_KEY_ID +// AWS_SECRET_ACCESS_KEY +// to be set along with the "bucket" and "region" arguments to be passed +package main + +import ( + "testing" + + "github.com/stretchr/testify/require" + + v1 "github.com/smithy-security/smithy/api/proto/v1" + + playwright "github.com/smithy-security/smithy/pkg/playwright/mock" + s3mock "github.com/smithy-security/smithy/pkg/s3/mock" + "github.com/smithy-security/smithy/pkg/testdata" +) + +func Test_run(t *testing.T) { + mockClient, err := playwright.NewMockClient() + require.NoError(t, err) + + pdfCalled := false + expected := []byte("this is a pdf") + mockClient.GetPDFOfPageCallBack = func(s1, s2 string) ([]byte, error) { + pdfCalled = true + return expected, nil + } + + mockS3Client, err := s3mock.NewMockClient("") + require.NoError(t, err) + s3Called := false + mockS3Client.UpsertCallback = func(s1, s2 string, b []byte) error { + s3Called = true + return nil + } + + err = run([]v1.EnrichedLaunchToolResponse{testdata.EnrichedLaunchToolResponse}, mockClient, mockS3Client) + require.NoError(t, err) + require.True(t, pdfCalled) + require.True(t, s3Called) + +} + +func Test_buildPdf(t *testing.T) { + mockClient, err := playwright.NewMockClient() + require.NoError(t, err) + + called := false + expected := []byte("this is a pdf") + mockClient.GetPDFOfPageCallBack = func(s1, s2 string) ([]byte, error) { + called = true + return expected, nil + } + _, result, err := buildPdf([]v1.EnrichedLaunchToolResponse{testdata.EnrichedLaunchToolResponse}, mockClient) + require.NoError(t, err) + require.Equal(t, called, true) + require.Equal(t, result, expected) +} diff --git a/pkg/playwright/mock/wrapper.go b/pkg/playwright/mock/wrapper.go new file mode 100644 index 000000000..8a644e16b --- /dev/null +++ b/pkg/playwright/mock/wrapper.go @@ -0,0 +1,19 @@ +package mock + +type MockClient struct { + StopCallBack func() error + GetPDFOfPageCallBack func(string, string) ([]byte, error) +} + +// NewClient returns an actual github client +func NewMockClient() (MockClient, error) { + return MockClient{}, nil +} + +func (c MockClient) Stop() error { + return c.StopCallBack() +} + +func (c MockClient) GetPDFOfPage(page, storePath string) ([]byte, error) { + return c.GetPDFOfPageCallBack(page, storePath) +} diff --git a/pkg/playwright/wrapper.go b/pkg/playwright/wrapper.go new file mode 100644 index 000000000..c57911251 --- /dev/null +++ b/pkg/playwright/wrapper.go @@ -0,0 +1,58 @@ +package wrapper + +import ( + "github.com/go-errors/errors" + "github.com/playwright-community/playwright-go" +) + +// Wrapper is the wrapper interface that allows the playwright client to be pluggable +type Wrapper interface { + Stop() error + GetPDFOfPage(string, string) ([]byte, error) +} + +// Client is the wrapper around google's go-github client +type Client struct { + playwright *playwright.Playwright +} + +// NewClient returns an actual github client +func NewClient() (Client, error) { + pw, err := playwright.Run() + if err != nil { + return Client{}, err + } + // create new playwright client + return Client{ + playwright: pw, + }, nil +} + +func (c Client) Stop() error { + return c.playwright.Stop() +} + +func (c Client) GetPDFOfPage(page, storePath string) ([]byte, error) { + browser, err := c.playwright.Chromium.Launch() + if err != nil { + return nil, err + } + + currentContext, err := browser.NewContext() + if err != nil { + return nil, err + } + + newPage, err := currentContext.NewPage() + if err != nil { + return nil, err + } + + if _, err = newPage.Goto(page); err != nil { + return nil, errors.Errorf("could not goto page %s in the browser: %w", page, err) + } + + return newPage.PDF(playwright.PagePdfOptions{ + Path: playwright.String(storePath), + }) +} diff --git a/pkg/s3/mock/wrapper.go b/pkg/s3/mock/wrapper.go new file mode 100644 index 000000000..868597fd4 --- /dev/null +++ b/pkg/s3/mock/wrapper.go @@ -0,0 +1,16 @@ +package wrapper + +// Client is the wrapper around s3 sdk +type Client struct { + UpsertCallback func(string, string, []byte) error +} + +// NewClient returns a client +func NewMockClient(region string) (Client, error) { + // create new playwright client + return Client{}, nil +} + +func (c Client) UpsertFile(filename, bucket string, pdfBytes []byte) error { + return c.UpsertCallback(filename, bucket, pdfBytes) +} diff --git a/pkg/s3/wrapper.go b/pkg/s3/wrapper.go new file mode 100644 index 000000000..3eb7538be --- /dev/null +++ b/pkg/s3/wrapper.go @@ -0,0 +1,65 @@ +package wrapper + +import ( + "bytes" + "log/slog" + "os" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/session" + "github.com/aws/aws-sdk-go/service/s3/s3manager" + "github.com/go-errors/errors" +) + +// Wrapper is the wrapper interface that allows the playwright client to be pluggable +type Wrapper interface { + // UpsertFile inserts or replaces a file in s3 if it doesn't exist + UpsertFile(string, string, []byte) error +} + +// Client is the wrapper around google's go-github client +type Client struct { + session *session.Session +} + +// NewClient returns a client +func NewClient(region string) (Client, error) { + sess, err := session.NewSession(&aws.Config{Region: aws.String(region)}) + if err != nil { + return Client{}, errors.Errorf("unable to start session with AWS API: %w", err) + } + // create new playwright client + return Client{ + session: sess, + }, nil +} + +func (c Client) UpsertFile(filename, bucket string, pdfBytes []byte) error { + //#nosec:G304 + data, err := os.ReadFile(filename) //#nosec:G304 + if err != nil { + return errors.Errorf("could not open file: %w", err) + } + + uploader := s3manager.NewUploader(c.session) + _, err = uploader.Upload(&s3manager.UploadInput{ + Bucket: aws.String(bucket), + Key: aws.String(filename), + Body: bytes.NewReader(data), + }) + if err != nil { + return errors.Errorf("unable to upload %s to %s: %w", filename, bucket, err) + } + + _, err = uploader.Upload(&s3manager.UploadInput{ + Bucket: aws.String(bucket), + Key: aws.String(filename), + Body: bytes.NewReader(pdfBytes), + }) + if err != nil { + return errors.Errorf("unable to upload %s to %s: %w", filename, bucket, err) + } + + slog.Info("uploaded", "filename", filename, "to", "bucket", bucket, "successfully") + return nil +}