From 6e8122779cf3d171783a59593a47da21f69bff8d Mon Sep 17 00:00:00 2001
From: sg <spyros@ocurity.com>
Date: Thu, 24 Oct 2024 16:21:56 +0100
Subject: [PATCH] add a pluggable aws-s3 wrapper

---
 pkg/s3/mock/wrapper.go | 16 ++++++++++
 pkg/s3/wrapper.go      | 67 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 pkg/s3/mock/wrapper.go
 create mode 100644 pkg/s3/wrapper.go

diff --git a/pkg/s3/mock/wrapper.go b/pkg/s3/mock/wrapper.go
new file mode 100644
index 000000000..eb9b549b9
--- /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..580897ac0
--- /dev/null
+++ b/pkg/s3/wrapper.go
@@ -0,0 +1,67 @@
+package wrapper
+
+import (
+	"bytes"
+	"fmt"
+	"log/slog"
+	"os"
+	"strings"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
+)
+
+// 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, 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{}, fmt.Errorf("unable to start session with AWS API: %w", err)
+	}
+	// create new playwright client
+	return Client{
+		session: sess,
+	}, nil
+}
+
+func (c Client) UpsertFile(htmlFilename, bucket, filenamePostfix string, pdfBytes []byte) error {
+	//#nosec:G304
+	data, err := os.ReadFile(htmlFilename) //#nosec:G304
+	if err != nil {
+		return fmt.Errorf("could not open file: %w", err)
+	}
+
+	uploader := s3manager.NewUploader(c.session)
+	_, err = uploader.Upload(&s3manager.UploadInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(htmlFilename + filenamePostfix),
+		Body:   bytes.NewReader(data),
+	})
+	if err != nil {
+		return fmt.Errorf("unable to upload %s to %s: %w", htmlFilename, bucket, err)
+	}
+	slog.Info("uploaded", "filename", htmlFilename, "to", "bucket", bucket, "successfully")
+
+	pdfFilename := strings.Replace(htmlFilename, ".html", "", -1) + ".pdf"
+	_, err = uploader.Upload(&s3manager.UploadInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(pdfFilename + filenamePostfix),
+		Body:   bytes.NewReader(pdfBytes),
+	})
+	if err != nil {
+		return fmt.Errorf("unable to upload %s to %s: %w", string(pdfBytes), bucket, err)
+	}
+	slog.Info("uploaded successfully", slog.String("filename", pdfFilename+filenamePostfix), slog.String("bucket", bucket))
+	return nil
+}