feat(teambition): support s3 upload method (close #4365)

AlistGo · May 13, 2023 · 3c4c2ad · 3c4c2ad
1 parent 077a525
commit 3c4c2ad
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 16 deletions.
diff --git a/drivers/teambition/driver.go b/drivers/teambition/driver.go
@@ -125,6 +125,9 @@ func (d *Teambition) Remove(ctx context.Context, obj model.Obj) error {
 }
 
 func (d *Teambition) Put(ctx context.Context, dstDir model.Obj, stream model.FileStreamer, up driver.UpdateProgress) error {
+	if d.UseS3UploadMethod {
+		return d.newUpload(ctx, dstDir, stream, up)
+	}
 	res, err := d.request("/api/v2/users/me", http.MethodGet, nil, nil)
 	if err != nil {
 		return err

diff --git a/drivers/teambition/meta.go b/drivers/teambition/meta.go
@@ -10,8 +10,9 @@ type Addition struct {
 	Cookie    string `json:"cookie" required:"true"`
 	ProjectID string `json:"project_id" required:"true"`
 	driver.RootID
-	OrderBy        string `json:"order_by" type:"select" options:"fileName,fileSize,updated,created" default:"fileName"`
-	OrderDirection string `json:"order_direction" type:"select" options:"Asc,Desc" default:"Asc"`
+	OrderBy           string `json:"order_by" type:"select" options:"fileName,fileSize,updated,created" default:"fileName"`
+	OrderDirection    string `json:"order_direction" type:"select" options:"Asc,Desc" default:"Asc"`
+	UseS3UploadMethod bool   `json:"use_s3_upload_method" default:"true"`
 }
 
 var config = driver.Config{

diff --git a/drivers/teambition/types.go b/drivers/teambition/types.go
@@ -66,3 +66,24 @@ type ChunkUpload struct {
 	PreviewExt     string      `json:"previewExt"`
 	LastUploadTime interface{} `json:"lastUploadTime"`
 }
+
+type UploadToken struct {
+	Sdk struct {
+		Endpoint         string `json:"endpoint"`
+		Region           string `json:"region"`
+		S3ForcePathStyle bool   `json:"s3ForcePathStyle"`
+		Credentials      struct {
+			AccessKeyId     string `json:"accessKeyId"`
+			SecretAccessKey string `json:"secretAccessKey"`
+			SessionToken    string `json:"sessionToken"`
+		} `json:"credentials"`
+	} `json:"sdk"`
+	Upload struct {
+		Bucket             string `json:"Bucket"`
+		Key                string `json:"Key"`
+		ContentDisposition string `json:"ContentDisposition"`
+		ContentType        string `json:"ContentType"`
+	} `json:"upload"`
+	Token       string `json:"token"`
+	DownloadUrl string `json:"downloadUrl"`
+}
diff --git a/drivers/teambition/util.go b/drivers/teambition/util.go
@@ -7,13 +7,16 @@ import (
 	"io"
 	"net/http"
 	"strconv"
-	"strings"
 	"time"
 
 	"github.com/alist-org/alist/v3/drivers/base"
 	"github.com/alist-org/alist/v3/internal/driver"
 	"github.com/alist-org/alist/v3/internal/model"
 	"github.com/alist-org/alist/v3/pkg/utils"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
 	"github.com/go-resty/resty/v2"
 	log "github.com/sirupsen/logrus"
 )
@@ -210,17 +213,56 @@ func (d *Teambition) finishUpload(file *FileUpload, parentId string) error {
 	return err
 }
 
-func getBetweenStr(str, start, end string) string {
-	n := strings.Index(str, start)
-	if n == -1 {
-		return ""
-	}
-	n = n + len(start)
-	str = string([]byte(str)[n:])
-	m := strings.Index(str, end)
-	if m == -1 {
-		return ""
-	}
-	str = string([]byte(str)[:m])
-	return str
+func (d *Teambition) newUpload(ctx context.Context, dstDir model.Obj, stream model.FileStreamer, up driver.UpdateProgress) error {
+	var uploadToken UploadToken
+	_, err := d.request("/api/awos/upload-token", http.MethodPost, func(req *resty.Request) {
+		req.SetBody(base.Json{
+			"category": "work",
+			"fileName": stream.GetName(),
+			"fileSize": stream.GetSize(),
+			"fileType": stream.GetMimetype(),
+			"payload": base.Json{
+				"involveMembers": []struct{}{},
+				"visible":        "members",
+			},
+			"scope": "project:" + d.ProjectID,
+		})
+	}, &uploadToken)
+	if err != nil {
+		return err
+	}
+	cfg := &aws.Config{
+		Credentials: credentials.NewStaticCredentials(
+			uploadToken.Sdk.Credentials.AccessKeyId, uploadToken.Sdk.Credentials.SecretAccessKey, uploadToken.Sdk.Credentials.SessionToken),
+		Region:           &uploadToken.Sdk.Region,
+		Endpoint:         &uploadToken.Sdk.Endpoint,
+		S3ForcePathStyle: &uploadToken.Sdk.S3ForcePathStyle,
+	}
+	ss, err := session.NewSession(cfg)
+	if err != nil {
+		return err
+	}
+	uploader := s3manager.NewUploader(ss)
+	input := &s3manager.UploadInput{
+		Bucket:             &uploadToken.Upload.Bucket,
+		Key:                &uploadToken.Upload.Key,
+		ContentDisposition: &uploadToken.Upload.ContentDisposition,
+		ContentType:        &uploadToken.Upload.ContentType,
+		Body:               stream,
+	}
+	_, err = uploader.UploadWithContext(ctx, input)
+	if err != nil {
+		return err
+	}
+	// finish upload
+	_, err = d.request("/api/works", http.MethodPost, func(req *resty.Request) {
+		req.SetBody(base.Json{
+			"fileTokens":     []string{uploadToken.Token},
+			"involveMembers": []struct{}{},
+			"visible":        "members",
+			"works":          []struct{}{},
+			"_parentId":      dstDir.GetID(),
+		})
+	}, nil)
+	return err
 }