s3.go

package s3

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net/url"
	"path"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"gitlab.com/uafrica/go-utils/errors"

	"github.com/aws/aws-sdk-go/aws/session"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/google/uuid"
)

// S3UploadResponse defines the structure of a standard JSON response to a PDF/CSV/etc request.
type S3UploadResponse struct {
	URL      string `json:"url"`
	Filename string `json:"filename"`
	Bucket   string `json:"bucket"`
	FileSize int    `json:"file_size"`
}

type MIMEType string

const (
	// MIMETypePDF defines the constant for the PDF MIME type.
	MIMETypePDF MIMEType = "application/pdf"

	// MIMETypeCSV defines the constant for the CSV MIME type.
	MIMETypeCSV MIMEType = "text/csv"

	// MIMETypeZIP defines the constant for the ZIP MIME type.
	MIMETypeZIP MIMEType = "application/zip"

	// MIMETypeJSON defines the constant for the JSON MIME type.
	MIMETypeJSON MIMEType = "application/json"

	// MIMETypeText defines the constant for the Plain text MIME type.
	MIMETypeText MIMEType = "text/plain"

	// MIMETypeImage defines the constant for the Image MIME type.
	MIMETypeImage MIMEType = "image/*"

	// MIMETypeDefault defines the constant for the default MIME type.
	MIMETypeDefault MIMEType = "application/octet-stream"

	// TypeXLS defines the constant for the XLS MIME type.
	MIMETypeXLS MIMEType = "application/vnd.ms-excel"

	// TypeXLSX defines the constant for the XLSX MIME type.
	MIMETypeXLSX MIMEType = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
)

type SessionWithHelpers struct {
	S3Session *s3.S3
}

func NewSession(session *session.Session) *SessionWithHelpers {
	return &SessionWithHelpers{
		S3Session: s3.New(session),
	}
}

func (s SessionWithHelpers) Upload(data []byte, bucket, fileName string, metadata *map[string]*string, isDebug bool) (*s3.PutObjectOutput, error) {
	mimeType := getTypeForFilename(fileName)
	putInput := &s3.PutObjectInput{
		Bucket:      aws.String(bucket),
		Key:         aws.String(fileName),
		ContentType: aws.String(string(mimeType)),
		Body:        bytes.NewReader(data),
	}

	if metadata != nil {
		putInput.Metadata = *metadata
	}

	response, err := s.S3Session.PutObject(putInput)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (s SessionWithHelpers) UploadWith1DayExpiry(data []byte, bucket, fileName string, mimeType MIMEType, isDebug bool) (string, error) {
	if mimeType == "" {
		mimeType = getTypeForFilename(fileName)
	}

	expiry := time.Now().Add(24 * time.Hour)
	putInput := &s3.PutObjectInput{
		Bucket:      aws.String(bucket),
		Key:         aws.String(fileName),
		ContentType: aws.String(string(mimeType)),
		Body:        bytes.NewReader(data),
		Expires:     &expiry,
	}

	_, err := s.S3Session.PutObject(putInput)
	if err != nil {
		return "", err
	}

	return s.GetSignedDownloadURL(bucket, fileName, 24*time.Hour, isDebug)
}

func (s SessionWithHelpers) UploadWithExpiry(data []byte, bucket, fileName string, expiryDuration time.Duration, mimeType MIMEType, isDebug bool) (string, error) {
	if mimeType == "" {
		mimeType = getTypeForFilename(fileName)
	}

	expiry := time.Now().Add(expiryDuration)
	putInput := &s3.PutObjectInput{
		Bucket:      aws.String(bucket),
		Key:         aws.String(fileName),
		ContentType: aws.String(string(mimeType)),
		Body:        bytes.NewReader(data),
		Expires:     &expiry,
	}

	_, err := s.S3Session.PutObject(putInput)
	if err != nil {
		return "", err
	}

	return s.GetSignedDownloadURL(bucket, fileName, 24*time.Hour, isDebug)
}

// UploadTempFile upload a file to S3 that will be automatically deleted after the expireDate
func (s SessionWithHelpers) UploadTempFile(data []byte, bucket, fileName string, metadata *map[string]*string, expireDate *time.Time) (*s3.PutObjectOutput, error) {
	mimeType := getTypeForFilename(fileName)
	putInput := &s3.PutObjectInput{
		Bucket:      aws.String(bucket),
		Key:         aws.String(fileName),
		ContentType: aws.String(string(mimeType)),
		Body:        bytes.NewReader(data),
		Expires:     expireDate,
	}

	if metadata != nil {
		putInput.Metadata = *metadata
	}

	response, err := s.S3Session.PutObject(putInput)
	if err != nil {
		return nil, err
	}

	return response, nil
}

// GetSignedDownloadURL gets a signed download URL for the duration. If scv is nil, a new session will be created.
func (s SessionWithHelpers) GetSignedDownloadURL(bucket string, fileName string, duration time.Duration, isDebug bool) (string, error) {
	getInput := &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(fileName),
	}
	getRequest, _ := s.S3Session.GetObjectRequest(getInput)

	fileExists, err := s.FileExists(bucket, fileName)
	if err != nil {
		return "", err
	}

	if !fileExists {
		return "", errors.Error("File does not exist")
	}

	return getRequest.Presign(duration)
}

func (s SessionWithHelpers) FileExists(bucket string, fileName string) (bool, error) {
	_, err := s.S3Session.HeadObject(&s3.HeadObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(fileName),
	})
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			switch aerr.Code() {
			case "NotFound": // s3.ErrCodeNoSuchKey does not work, aws is missing this error code so we hardwire a string
				return false, nil
			default:
				return false, err
			}
		}
		return false, err
	}
	return true, nil
}

// UploadWithFileExtension will upload a file to S3 and return a standard S3UploadResponse.
func (s SessionWithHelpers) UploadWithFileExtension(data []byte, bucket, filePrefix string, fileExt string, mimeType MIMEType, isDebug bool) (*S3UploadResponse, error) {
	fileName := fmt.Sprintf("%s_%s.%s", filePrefix, uuid.New().String(), fileExt)

	uploadUrl, err := s.UploadWith1DayExpiry(data, bucket, fileName, mimeType, isDebug)
	if err != nil {
		return nil, err
	}

	fileSizeInBytes := binary.Size(data)

	response := &S3UploadResponse{
		URL:      uploadUrl,
		Filename: fileName,
		Bucket:   bucket,
		FileSize: fileSizeInBytes,
	}

	return response, nil
}

func getTypeForFilename(f string) MIMEType {
	ext := strings.ToLower(path.Ext(f))

	switch ext {
	case "pdf":
		return MIMETypePDF
	case "csv":
		return MIMETypeCSV
	case "zip":
		return MIMETypeZIP
	case "txt":
		return MIMETypeText
	}

	return MIMETypeDefault
}

func (s SessionWithHelpers) GetObject(bucket string, fileName string, isDebug bool) (*s3.GetObjectOutput, error) {
	getInput := &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(fileName),
	}
	getObjectOutput, err := s.S3Session.GetObject(getInput)
	if err != nil {
		return nil, err
	}
	return getObjectOutput, nil
}

func (s SessionWithHelpers) GetObjectMetadata(bucket string, fileName string, isDebug bool) (map[string]*string, error) {
	headObjectInput := &s3.HeadObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(fileName),
	}
	headObjectOutput, err := s.S3Session.HeadObject(headObjectInput)
	if err != nil {
		return nil, err
	}
	return headObjectOutput.Metadata, nil
}

// MoveObjectBucketToBucket - Move object from one S3 bucket to another
func (s SessionWithHelpers) MoveObjectBucketToBucket(sourceBucket string, destinationBucket string, sourceFileName string, destinationFileName string, isDebug bool) error {

	err := s.CopyObjectBucketToBucket(sourceBucket, destinationBucket, sourceFileName, destinationFileName, isDebug)
	if err != nil {
		return err
	}

	err = s.DeleteObjectFromBucket(sourceBucket, sourceFileName, isDebug)
	if err != nil {
		return err
	}

	return nil
}

// CopyObjectBucketToBucket - Copy an object from one S3 bucket to another
func (s SessionWithHelpers) CopyObjectBucketToBucket(sourceBucket string, destinationBucket string, sourceFileName string, destinationFilename string, isDebug bool) error {
	// copy the file
	copySource := url.QueryEscape(sourceBucket + "/" + sourceFileName)
	copyObjectInput := &s3.CopyObjectInput{
		Bucket:     aws.String(destinationBucket),   //destination bucket
		CopySource: aws.String(copySource),          //source path (ie: myBucket/myFile.csv)
		Key:        aws.String(destinationFilename), //filename on destination
	}
	_, err := s.S3Session.CopyObject(copyObjectInput)
	if err != nil {
		return err
	}

	// wait to see if the file copied successfully
	err = s.S3Session.WaitUntilObjectExists(&s3.HeadObjectInput{Bucket: aws.String(destinationBucket), Key: aws.String(destinationFilename)})
	if err != nil {
		return err
	}

	return nil
}

// DeleteObjectFromBucket - Delete an object from an S3 bucket
func (s SessionWithHelpers) DeleteObjectFromBucket(bucket string, fileName string, isDebug bool) error {
	// delete the file
	deleteObjectInput := &s3.DeleteObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(fileName),
	}
	_, err := s.S3Session.DeleteObject(deleteObjectInput)
	if err != nil {
		return err
	}

	// wait to see if the file deleted successfully
	err = s.S3Session.WaitUntilObjectNotExists(&s3.HeadObjectInput{
		Bucket: aws.String(bucket),   // the bucket we are deleting from
		Key:    aws.String(fileName), // the filename we are deleting
	})
	if err != nil {
		return err
	}

	return nil
}

func GetS3FileKey(fileName string, folder string) string {
	// Trim leading and trailing slashes
	fileName = strings.TrimLeft(fileName, "/")
	fileName = strings.TrimRight(fileName, "/")

	folder = strings.TrimLeft(folder, "/")
	folder = strings.TrimRight(folder, "/")

	return "/" + folder + "/" + fileName
}