WIP: Support Azure and GCP

404handler.go

@@ -8,9 +8,9 @@ import (
 func Handle404(helmet SimpleHelmet) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		for key, val := range helmet.headers {
-			w.Header().Set(key,val)
+			w.Header().Set(key, val)
 		}
 		w.Header().Set("Content-Type", "application/json")
 		w.WriteHeader(http.StatusNotFound)
 	})
-}
+}

abs.go

@@ -0,0 +1,149 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container"
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/blob"
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/bloberror"
+	"path"
+	"log"
+	"os"
+)
+
+type ABS_Manager struct {
+	BucketMgr
+}
+
+func (self *ABS_Manager) handleError(err error) {
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+}
+
+func (self *ABS_Manager) getServiceClient() *azblob.Client {
+	// Create a new service client with token credential
+	accountName, ok := os.LookupEnv("AZURE_STORAGE_ACCOUNT_NAME")
+	if !ok {
+		panic("AZURE_STORAGE_ACCOUNT_NAME could not be found")
+	}
+
+	serviceURL := fmt.Sprintf("https://%s.blob.core.windows.net/", accountName)
+
+	credential, err := azidentity.NewDefaultAzureCredential(nil)
+	self.handleError(err)
+
+	client, err := azblob.NewClient(serviceURL, credential, nil)
+	self.handleError(err)
+	return client
+}
+
+func (self *ABS_Manager) getContainerClient(containerName string) *container.Client {
+	accountName, ok := os.LookupEnv("AZURE_STORAGE_ACCOUNT_NAME")
+	if !ok {
+		panic("AZURE_STORAGE_ACCOUNT_NAME could not be found")
+	}
+
+	containerURL := fmt.Sprintf("https://%s.blob.core.windows.net/%s", accountName, containerName)
+
+	cred, err := azidentity.NewDefaultAzureCredential(nil)
+	self.handleError(err)
+
+	containerClient, err := container.NewClient(containerURL, cred, nil)
+	self.handleError(err)
+	return containerClient
+}
+
+func (self *ABS_Manager) getBlobClient(containerName string, blobName string) *blob.Client {
+	// From the Azure portal, get your Storage account blob service URL endpoint.
+	accountName, accountKey := os.Getenv("AZURE_STORAGE_ACCOUNT_NAME"), os.Getenv("AZURE_STORAGE_ACCOUNT_KEY")
+
+	blobURL := fmt.Sprintf("https://%s.blob.core.windows.net/%s/%s", accountName, containerName, blobName)
+	credential, err := azblob.NewSharedKeyCredential(accountName, accountKey)
+	self.handleError(err)
+	blobClient, err := blob.NewClientWithSharedKeyCredential(blobURL, credential, nil)
+	self.handleError(err)
+	return blobClient
+}
+
+func (self *ABS_Manager) listBuckets() []string {
+
+	client := self.getServiceClient()
+
+	pager := client.NewListContainersPager(&azblob.ListContainersOptions{
+		Include: azblob.ListContainersInclude{Metadata: true, Deleted: false},
+	})
+
+	var buckets []string
+
+	for pager.More() {
+		resp, err := pager.NextPage(context.TODO())
+		self.handleError(err) // if err is not nil, break the loop.
+		for _, _container := range resp.ContainerItems {
+			buckets = append(buckets, *_container.Name)
+		}
+	}
+	return buckets
+}
+
+func (self *ABS_Manager) bucketExists(bucket string) (bool, error) {
+	client := self.getContainerClient(bucket)
+	_, err := client.GetProperties(context.TODO(), nil)
+
+	if bloberror.HasCode(err, bloberror.ContainerNotFound) {
+		return false, err
+	} else {
+		return true, nil
+	}
+}
+
+func (self *ABS_Manager) keyExists(bucket string, key string) (bool, error) {
+	client := self.getBlobClient(bucket, key)
+	_, err := client.GetProperties(context.TODO(), nil)
+
+	if bloberror.HasCode(err, bloberror.BlobNotFound) {
+		return false, err
+	} else {
+		return true, nil
+	}
+}
+
+func (self *ABS_Manager) readFile(bucket string, item string) ([]byte, error) {
+
+	client := self.getServiceClient()
+	// Download the blob
+	downloadResponse, err := client.DownloadStream(context.TODO(), bucket, item, nil)
+	self.handleError(err)
+
+	downloadedData := bytes.Buffer{}
+	retryReader := downloadResponse.NewRetryReader(context.TODO(), &azblob.RetryReaderOptions{})
+	_, err = downloadedData.ReadFrom(retryReader)
+	self.handleError(err)
+
+	err = retryReader.Close()
+	self.handleError(err)
+
+	return downloadedData.Bytes(), nil
+}
+
+func (self *ABS_Manager) copyFile(bucket string, item string, other string) error {
+
+	data, _ := self.readFile(bucket, item)
+
+	client := self.getServiceClient()
+
+	_, err := client.UploadBuffer(context.TODO(), path.Dir(other), path.Base(other), data, &azblob.UploadBufferOptions{})
+	self.handleError(err)
+	return err
+}
+
+func (self *ABS_Manager) deleteFile(bucket string, item string) error {
+	client := self.getServiceClient()
+	// Delete the blob.
+	_, err := client.DeleteBlob(context.TODO(), bucket, item, nil)
+	self.handleError(err)
+	return err
+}

bucket_interface.go

@@ -0,0 +1,10 @@
+package main
+
+// Defining an interface
+type BucketInterface interface {
+	bucketExists(bucket string) (bool, error)
+	keyExists(bucket string, key string) (bool, error)
+	readFile(bucket string, item string) ([]byte, error)
+	copyFile(bucket string, item string, other string) error
+	deleteFile(bucket string, item string) error
+}

bucket_interface_manager.go

@@ -0,0 +1,213 @@
+package main
+
+import (
+	// standard
+	"encoding/json"
+	"errors"
+	"log"
+	"net/http"
+)
+
+func _getQurantineFilesBucket(qurantineFilesBucket string) string {
+	// input has more priority
+	if qurantineFilesBucket != "" {
+		return qurantineFilesBucket
+	}
+	if quarantine_files_bucket != "" {
+		return quarantine_files_bucket
+	}
+	return ""
+}
+
+func _getCleanFilesBucket(cleanFilesBucket string) string {
+	// input has more priority
+	if cleanFilesBucket != "" {
+		return cleanFilesBucket
+	}
+	if clean_files_bucket != "" {
+		return clean_files_bucket
+	}
+	return ""
+}
+
+func validateInputBucket(w http.ResponseWriter, bucket string, bucketInterface BucketInterface) error {
+	if bucket == "" {
+		errorResponse(w, "Invalid input bucket", http.StatusUnprocessableEntity)
+		return errors.New("Invalid input bucket")
+	}
+
+	bucketExists, err := bucketInterface.bucketExists(bucket)
+
+	if err != nil {
+		errorResponse(w, err.Error(), http.StatusInternalServerError)
+		return err
+	}
+	if !bucketExists {
+		errorResponse(w, "Bucket: "+bucket+" does not exists", http.StatusUnprocessableEntity)
+		return errors.New("Bucket: " + bucket + " does not exists")
+	}
+	return nil
+}
+
+func validateInputKey(w http.ResponseWriter, bucket string, key string, bucketInterface BucketInterface) error {
+	if key == "" {
+		errorResponse(w, "Invalid input key", http.StatusUnprocessableEntity)
+		return errors.New("Invalid input key")
+	}
+
+	keyExists, err := bucketInterface.keyExists(bucket, key)
+	if err != nil {
+		errorResponse(w, err.Error(), http.StatusInternalServerError)
+		return err
+	}
+	if !keyExists {
+		errorResponse(w, "Key: "+key+" does not exist in Bucket: "+bucket, http.StatusUnprocessableEntity)
+		return errors.New("Key: " + key + " does not exist in Bucket: " + bucket)
+	}
+	return nil
+}
+
+func validateQrantineFilesBucket(w http.ResponseWriter, qurantineFilesBucket string, bucketInterface BucketInterface) error {
+	var bucket = _getQurantineFilesBucket(qurantineFilesBucket)
+
+	if bucket == "" {
+		errorResponse(w, "Invalid qurantine files bucket", http.StatusBadRequest)
+		return errors.New("Invalid qurantine files bucket")
+
+	} else {
+		err := validateInputBucket(w, bucket, bucketInterface)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func validateCleanFilesBucket(w http.ResponseWriter, cleanFilesBucket string, bucketInterface BucketInterface) error {
+
+	var bucket = _getCleanFilesBucket(cleanFilesBucket)
+
+	if bucket == "" {
+		errorResponse(w, "Invalid clean files bucket", http.StatusBadRequest)
+		return errors.New("Invalid clean files bucket")
+
+	} else {
+		err := validateInputBucket(w, bucket, bucketInterface)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+
+}
+
+func validateInputData(w http.ResponseWriter, data *ScanObject, bucketInterface BucketInterface) error {
+
+	err := validateInputBucket(w, data.BucketName, bucketInterface)
+	if err != nil {
+		return err
+	}
+
+	err = validateInputKey(w, data.BucketName, data.Key, bucketInterface)
+	if err != nil {
+		return err
+	}
+
+	err = validateQrantineFilesBucket(w, data.QurantineFilesBucket, bucketInterface)
+	if err != nil {
+		return err
+	}
+
+	err = validateCleanFilesBucket(w, data.CleanFilesBucket, bucketInterface)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func ScanBucketObject(w http.ResponseWriter, r *http.Request, bucketInterface BucketInterface) {
+
+	data := new(ScanObject)
+	err := decodeJSONBody(w, r, &data)
+	if err != nil {
+		var mr *malformedRequest
+		if errors.As(err, &mr) {
+			errorResponse(w, mr.msg, mr.status)
+		} else {
+			log.Println(err.Error())
+			errorResponse(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
+		}
+		return
+	}
+
+	err = validateInputData(w, data, bucketInterface)
+	if err != nil {
+		elog.Println(" validateInputData failed " + err.Error())
+		return
+	}
+
+	resp, _ := json.Marshal(data)
+	info.Println(" Received ScanS3 request " + string(resp))
+
+	byteData, err := bucketInterface.readFile(data.BucketName, data.Key)
+	if err != nil {
+		elog.Println(err)
+		errorResponse(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
+		return
+	}
+
+	// send request for scanning
+	newRequest := NewScanStreamRequest(byteData)
+	scanstreamrequests <- newRequest
+
+	response := <-newRequest.ResponseChan
+
+	err = response.err
+
+	if err != nil {
+		elog.Println(err)
+		errorResponse(w, err.Error(), http.StatusInternalServerError)
+		return
+	} else {
+		if response.data.Status == "INFECTED" {
+			elog.Println("Key " + data.Key + " from bucket " + data.BucketName + " is Infected")
+			err = bucketInterface.copyFile(data.BucketName, data.Key, _getQurantineFilesBucket(data.QurantineFilesBucket))
+			if err != nil {
+				elog.Println(err)
+				errorResponse(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+			err = bucketInterface.deleteFile(data.BucketName, data.Key)
+			if err != nil {
+				elog.Println(err)
+				errorResponse(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		} else if response.data.Status == "CLEAN" {
+			info.Println("Key " + data.Key + " from bucket " + data.BucketName + " is Clean")
+			err = bucketInterface.copyFile(data.BucketName, data.Key, _getCleanFilesBucket(data.CleanFilesBucket))
+			if err != nil {
+				elog.Println(err)
+				errorResponse(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+			err = bucketInterface.deleteFile(data.BucketName, data.Key)
+			if err != nil {
+				elog.Println(err)
+				errorResponse(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		}
+	}
+
+	output, err := json.Marshal(response.data)
+	if err != nil {
+		elog.Println(err)
+		errorResponse(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	sendJsonResponse(w, output)
+	//fmt.Fprintf(w, string(output))
+}

bucket_manager.go

@@ -0,0 +1,4 @@
+package main
+
+type BucketMgr struct {
+}