diff --git a/cmd/fsb/run.go b/cmd/fsb/run.go index b38a00b27..0dbcf7d1d 100644 --- a/cmd/fsb/run.go +++ b/cmd/fsb/run.go @@ -6,6 +6,7 @@ import ( "EverythingSuckz/fsb/internal/cache" "EverythingSuckz/fsb/internal/routes" "EverythingSuckz/fsb/internal/types" + "EverythingSuckz/fsb/internal/uploader" "EverythingSuckz/fsb/internal/utils" "fmt" "net/http" @@ -38,6 +39,7 @@ func runApp(cmd *cobra.Command, args []string) { if err != nil { log.Panic("Failed to start main bot", zap.Error(err)) } + uploader.Load(log, mainBot.Dispatcher) cache.InitCache(log) workers, err := bot.StartWorkers(log) if err != nil { diff --git a/config/config.go b/config/config.go index 94bd059e5..70021fbbd 100644 --- a/config/config.go +++ b/config/config.go @@ -50,6 +50,7 @@ type config struct { UserSession string `envconfig:"USER_SESSION"` UsePublicIP bool `envconfig:"USE_PUBLIC_IP" default:"false"` AllowedUsers allowedUsers `envconfig:"ALLOWED_USERS"` + Drive115Cookie string `envconfig:"DRIVE115_COOKIE"` MultiTokens []string } diff --git a/internal/routes/routes.go b/internal/routes/routes.go index 445862756..d7b53cfbb 100644 --- a/internal/routes/routes.go +++ b/internal/routes/routes.go @@ -28,6 +28,6 @@ func Load(log *zap.Logger, r *gin.Engine) { Type := reflect.TypeOf(&allRoutes{log}) Value := reflect.ValueOf(&allRoutes{log}) for i := 0; i < Type.NumMethod(); i++ { - Type.Method(i).Func.Call([]reflect.Value{Value, reflect.ValueOf(route)}) + Type.Method(i).Func.Call([]reflect.Value{Value, reflect.ValueOf(&Route{Engine: r})}) } } diff --git a/internal/routes/upload.go b/internal/routes/upload.go new file mode 100644 index 000000000..64324a012 --- /dev/null +++ b/internal/routes/upload.go @@ -0,0 +1,76 @@ +package routes + +import ( + "EverythingSuckz/fsb/config" + "EverythingSuckz/fsb/internal/bot" + "EverythingSuckz/fsb/internal/utils" + "EverythingSuckz/fsb/pkg/drive115" + "fmt" + "net/http" + "strconv" + + "github.com/gin-gonic/gin" +) + +func (e *allRoutes) LoadUpload(r *Route) { + log := e.log.Named("Upload") + defer log.Info("Loaded upload route") + r.Engine.GET("/upload/:messageID", getUploadRoute) +} + +func getUploadRoute(ctx *gin.Context) { + w := ctx.Writer + + messageIDParm := ctx.Param("messageID") + messageID, err := strconv.Atoi(messageIDParm) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + authHash := ctx.Query("hash") + if authHash == "" { + http.Error(w, "missing hash param", http.StatusBadRequest) + return + } + + worker := bot.GetNextWorker() + + file, err := utils.FileFromMessage(ctx, worker.Client, messageID) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + expectedHash := utils.PackFile( + file.FileName, + file.FileSize, + file.MimeType, + file.ID, + ) + if !utils.CheckHash(authHash, expectedHash) { + http.Error(w, "invalid hash", http.StatusBadRequest) + return + } + + driveClient := drive115.NewClient(config.ValueOf.Drive115Cookie) + uploadInfo, err := driveClient.GetUploadInfo() + if err != nil { + http.Error(w, fmt.Sprintf("Failed to get upload info: %s", err), http.StatusInternalServerError) + return + } + + reader, err := utils.NewTelegramReader(ctx, worker.Client, file.Location, 0, file.FileSize-1, file.FileSize) + if err != nil { + http.Error(w, fmt.Sprintf("Failed to create telegram reader: %s", err), http.StatusInternalServerError) + return + } + + result, err := driveClient.Upload(reader, uploadInfo.UploadURL) + if err != nil { + http.Error(w, fmt.Sprintf("Failed to upload file: %s", err), http.StatusInternalServerError) + return + } + + ctx.String(http.StatusOK, "File uploaded successfully. File ID: %s", result.FileID) +} diff --git a/internal/uploader/upload.go b/internal/uploader/upload.go new file mode 100644 index 000000000..dfe78574f --- /dev/null +++ b/internal/uploader/upload.go @@ -0,0 +1,59 @@ +package uploader + +import ( + "EverythingSuckz/fsb/config" + "EverythingSuckz/fsb/internal/bot" + "EverythingSuckz/fsb/internal/utils" + "EverythingSuckz/fsb/pkg/drive115" + "fmt" + + "github.com/celestix/gotgproto/dispatcher" + "github.com/celestix/gotgproto/dispatcher/handlers" + "github.com/celestix/gotgproto/ext" +) + +import "go.uber.org/zap" + +func Load(log *zap.Logger, dispatcher dispatcher.Dispatcher) { + log = log.Named("uploader") + defer log.Info("Initialized uploader command handlers") + dispatcher.AddHandler( + handlers.NewCommand("upload", upload), + ) +} + +func upload(ctx *ext.Context, u *ext.Update) error { + if u.EffectiveMessage.ReplyToMessage == nil { + ctx.Reply(u, "Reply to a message to upload.", nil) + return dispatcher.EndGroups + } + + driveClient := drive115.NewClient(config.ValueOf.Drive115Cookie) + uploadInfo, err := driveClient.GetUploadInfo() + if err != nil { + ctx.Reply(u, fmt.Sprintf("Failed to get upload info: %s", err), nil) + return dispatcher.EndGroups + } + + worker := bot.GetNextWorker() + file, err := utils.FileFromMessage(ctx, worker.Client, u.EffectiveMessage.ReplyToMessage.GetID()) + if err != nil { + ctx.Reply(u, fmt.Sprintf("Failed to get file from message: %s", err), nil) + return dispatcher.EndGroups + } + + reader, err := utils.NewTelegramReader(ctx, worker.Client, file.Location, 0, file.FileSize-1, file.FileSize) + if err != nil { + ctx.Reply(u, fmt.Sprintf("Failed to create telegram reader: %s", err), nil) + return dispatcher.EndGroups + } + + result, err := driveClient.Upload(reader, uploadInfo.UploadURL) + if err != nil { + ctx.Reply(u, fmt.Sprintf("Failed to upload file: %s", err), nil) + return dispatcher.EndGroups + } + + ctx.Reply(u, fmt.Sprintf("File uploaded successfully. File ID: %s", result.FileID), nil) + return dispatcher.EndGroups +} diff --git a/pkg/drive115/api.go b/pkg/drive115/api.go new file mode 100644 index 000000000..53d69a977 --- /dev/null +++ b/pkg/drive115/api.go @@ -0,0 +1,82 @@ +package drive115 + +import ( + "encoding/json" + "fmt" + "io" + "net/http" +) + +const ( + apiGetUploadInfo = "https://uplb.115.com/3.0/getuploadinfo.php" +) + +type Client struct { + cookie string + httpClient *http.Client +} + +func NewClient(cookie string) *Client { + return &Client{ + cookie: cookie, + httpClient: &http.Client{}, + } +} + +type UploadInfo struct { + UploadURL string `json:"upload_url"` +} + +func (c *Client) GetUploadInfo() (*UploadInfo, error) { + req, err := http.NewRequest("GET", apiGetUploadInfo, nil) + if err != nil { + return nil, err + } + req.Header.Set("Cookie", c.cookie) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("get upload info failed with status: %s", resp.Status) + } + + var info UploadInfo + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return nil, err + } + + return &info, nil +} + +type UploadResult struct { + FileID string `json:"file_id"` +} + +func (c *Client) Upload(reader io.Reader, uploadURL string) (*UploadResult, error) { + req, err := http.NewRequest("POST", uploadURL, reader) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/octet-stream") + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("upload failed with status: %s", resp.Status) + } + + var result UploadResult + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, err + } + + return &result, nil +} \ No newline at end of file