diff --git a/bot.go b/bot.go index 86bbfc8..9fe384f 100644 --- a/bot.go +++ b/bot.go @@ -15,7 +15,7 @@ import ( "go.mongodb.org/mongo-driver/v2/mongo" ) -type BotSettings struct { +type BotOpts struct { Token string Debug bool ErrorTemplate string @@ -24,10 +24,12 @@ type BotSettings struct { LoggerBasePath string UseRequestLogger bool WriteToFile bool + UseTestServer bool + APIUrl string } -func LoadSettingsFromEnv() *BotSettings { - return &BotSettings{ +func LoadOptsFromEnv() *BotOpts { + return &BotOpts{ Token: os.Getenv("TG_TOKEN"), Debug: os.Getenv("DEBUG") == "true", ErrorTemplate: os.Getenv("ERROR_TEMPLATE"), @@ -35,6 +37,8 @@ func LoadSettingsFromEnv() *BotSettings { UpdateTypes: strings.Split(os.Getenv("UPDATE_TYPES"), ";"), UseRequestLogger: os.Getenv("USE_REQ_LOG") == "true", WriteToFile: os.Getenv("WRITE_TO_FILE") == "true", + UseTestServer: os.Getenv("USE_TEST_SERVER") == "true", + APIUrl: os.Getenv("API_URL"), } } @@ -70,9 +74,10 @@ type Bot struct { updateQueue *extypes.Queue[*tgapi.Update] } -func NewBot(settings *BotSettings) *Bot { +func NewBot(settings *BotOpts) *Bot { updateQueue := extypes.CreateQueue[*tgapi.Update](256) - api := tgapi.NewAPI(settings.Token) + apiOpts := tgapi.NewAPIOpts(settings.Token).SetAPIUrl(settings.APIUrl).UseTestServer(settings.UseTestServer) + api := tgapi.NewAPI(apiOpts) bot := &Bot{ updateOffset: 0, plugins: make([]Plugin, 0), debug: settings.Debug, errorTemplate: "%s", prefixes: settings.Prefixes, updateTypes: make([]tgapi.UpdateType, 0), runners: make([]Runner, 0), @@ -237,7 +242,6 @@ func (b *Bot) Run() { return } - b.logger.Infoln("Executing runners...") b.ExecRunners() b.logger.Infoln("Bot running. Press CTRL+C to exit.") diff --git a/cmd_generator.go b/cmd_generator.go index 7d67a2c..b4f8c35 100644 --- a/cmd_generator.go +++ b/cmd_generator.go @@ -1,6 +1,7 @@ package laniakea import ( + "errors" "fmt" "strings" @@ -35,6 +36,8 @@ func generateBotCommandForPlugin(pl Plugin) []tgapi.BotCommand { return commands } +var ErrTooManyCommands = errors.New("too many commands. max 100") + func (b *Bot) AutoGenerateCommands() error { _, err := b.api.DeleteMyCommands(tgapi.DeleteMyCommandsP{}) if err != nil { @@ -45,6 +48,9 @@ func (b *Bot) AutoGenerateCommands() error { for _, pl := range b.plugins { commands = append(commands, generateBotCommandForPlugin(pl)...) } + if len(commands) > 100 { + return ErrTooManyCommands + } privateChatsScope := &tgapi.BotCommandScope{Type: tgapi.BotCommandScopePrivateType} groupChatsScope := &tgapi.BotCommandScope{Type: tgapi.BotCommandScopeGroupType} diff --git a/tgapi/api.go b/tgapi/api.go index fe8f5b4..cafa5b8 100644 --- a/tgapi/api.go +++ b/tgapi/api.go @@ -13,17 +13,49 @@ import ( "git.nix13.pw/scuroneko/slog" ) -type API struct { - token string - client *http.Client - Logger *slog.Logger +type APIOpts struct { + token string + client *http.Client + useTestServer bool + apiUrl string } -func NewAPI(token string) *API { +func NewAPIOpts(token string) *APIOpts { + return &APIOpts{token: token, client: nil, useTestServer: false, apiUrl: "https://api.telegram.org"} +} +func (opts *APIOpts) SetHTTPClient(client *http.Client) *APIOpts { + if client != nil { + opts.client = client + } + return opts +} +func (opts *APIOpts) UseTestServer(use bool) *APIOpts { + opts.useTestServer = use + return opts +} +func (opts *APIOpts) SetAPIUrl(apiUrl string) *APIOpts { + if apiUrl != "" { + opts.apiUrl = apiUrl + } + return opts +} + +type API struct { + token string + client *http.Client + Logger *slog.Logger + useTestServer bool + apiUrl string +} + +func NewAPI(opts *APIOpts) *API { l := slog.CreateLogger().Level(utils.GetLoggerLevel()).Prefix("API") l.AddWriter(l.CreateJsonStdoutWriter()) - client := &http.Client{Timeout: time.Second * 45} - return &API{token, client, l} + client := opts.client + if client == nil { + client = &http.Client{Timeout: time.Second * 45} + } + return &API{opts.token, client, l, opts.useTestServer, opts.apiUrl} } func (api *API) CloseApi() error { return api.Logger.Close() @@ -52,8 +84,12 @@ func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R, } buf := bytes.NewBuffer(data) - u := fmt.Sprintf("https://api.telegram.org/bot%s/%s", api.token, r.method) - req, err := http.NewRequestWithContext(ctx, "POST", u, buf) + methodPrefix := "" + if api.useTestServer { + methodPrefix = "/test" + } + url := fmt.Sprintf("%s/bot%s%s/%s", api.apiUrl, api.token, methodPrefix, r.method) + req, err := http.NewRequestWithContext(ctx, "POST", url, buf) if err != nil { return zero, err } @@ -61,15 +97,16 @@ func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R, req.Header.Set("Accept", "application/json") req.Header.Set("User-Agent", fmt.Sprintf("Laniakea/%s", utils.VersionString)) - api.Logger.Debugln("REQ", r.method, buf.String()) + api.Logger.Debugln("REQ", api.apiUrl, r.method, buf.String()) res, err := api.client.Do(req) if err != nil { return zero, err } - defer res.Body.Close() + defer func(Body io.ReadCloser) { + _ = Body.Close() + }(res.Body) - reader := io.LimitReader(res.Body, 10<<20) - data, err = io.ReadAll(reader) + data, err = readBody(res.Body) if err != nil { return zero, err } @@ -77,9 +114,21 @@ func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R, if res.StatusCode != http.StatusOK { return zero, fmt.Errorf("unexpected status code: %d, %s", res.StatusCode, string(data)) } + return parseBody[R](data) +} +func (r TelegramRequest[R, P]) Do(api *API) (R, error) { + return r.DoWithContext(context.Background(), api) +} + +func readBody(body io.ReadCloser) ([]byte, error) { + reader := io.LimitReader(body, 10<<20) + return io.ReadAll(reader) +} +func parseBody[R any](data []byte) (R, error) { + var zero R var resp ApiResponse[R] - err = json.Unmarshal(data, &resp) + err := json.Unmarshal(data, &resp) if err != nil { return zero, err } @@ -87,9 +136,4 @@ func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R, return zero, fmt.Errorf("[%d] %s", resp.ErrorCode, resp.Description) } return resp.Result, nil - -} -func (r TelegramRequest[R, P]) Do(api *API) (R, error) { - ctx := context.Background() - return r.DoWithContext(ctx, api) } diff --git a/tgapi/uploader_api.go b/tgapi/uploader_api.go index 64fe7ca..ec0f9d7 100644 --- a/tgapi/uploader_api.go +++ b/tgapi/uploader_api.go @@ -3,9 +3,7 @@ package tgapi import ( "bytes" "context" - "encoding/json" "fmt" - "io" "mime/multipart" "net/http" "path/filepath" @@ -66,39 +64,22 @@ func NewUploaderRequest[R, P any](method string, params P, files ...UploaderFile } func (u UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader) (R, error) { var zero R - url := fmt.Sprintf("https://api.telegram.org/bot%s/%s", up.api.token, u.method) - buf := bytes.NewBuffer(nil) - w := multipart.NewWriter(buf) - - for _, file := range u.files { - fw, err := w.CreateFormFile(string(file.field), file.filename) - if err != nil { - _ = w.Close() - return zero, err - } - _, err = fw.Write(file.data) - if err != nil { - _ = w.Close() - return zero, err - } - } - - err := utils.Encode(w, u.params) - if err != nil { - _ = w.Close() - return zero, err - } - err = w.Close() + buf, contentType, err := prepareMultipart(u.files, u.params) if err != nil { return zero, err } + methodPrefix := "" + if up.api.useTestServer { + methodPrefix = "/test" + } + url := fmt.Sprintf("%s/bot%s%s/%s", up.api.apiUrl, up.api.token, methodPrefix, u.method) req, err := http.NewRequestWithContext(ctx, "POST", url, buf) if err != nil { return zero, err } - req.Header.Set("Content-Type", w.FormDataContentType()) + req.Header.Set("Content-Type", contentType) req.Header.Set("Accept", "application/json") req.Header.Set("User-Agent", fmt.Sprintf("Laniakea/%s", utils.VersionString)) @@ -109,30 +90,45 @@ func (u UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader) } defer res.Body.Close() - reader := io.LimitReader(res.Body, 10<<20) - body, err := io.ReadAll(reader) - if err != nil { - return zero, err - } + body, err := readBody(res.Body) up.logger.Debugln("UPLOADER RES", u.method, string(body)) if res.StatusCode != http.StatusOK { return zero, fmt.Errorf("unexpected status code: %d, %s", res.StatusCode, string(body)) } - var resp ApiResponse[R] - err = json.Unmarshal(body, &resp) - if err != nil { - return zero, err - } - if !resp.Ok { - return zero, fmt.Errorf("[%d] %s", resp.ErrorCode, resp.Description) - } - return resp.Result, nil + return parseBody[R](body) } func (u UploaderRequest[R, P]) Do(up *Uploader) (R, error) { return u.DoWithContext(context.Background(), up) } +func prepareMultipart[P any](files []UploaderFile, params P) (*bytes.Buffer, string, error) { + buf := bytes.NewBuffer(nil) + w := multipart.NewWriter(buf) + + for _, file := range files { + fw, err := w.CreateFormFile(string(file.field), file.filename) + if err != nil { + _ = w.Close() + return buf, w.FormDataContentType(), err + } + + _, err = fw.Write(file.data) + if err != nil { + _ = w.Close() + return buf, w.FormDataContentType(), err + } + } + + err := utils.Encode(w, params) + if err != nil { + _ = w.Close() + return buf, w.FormDataContentType(), err + } + err = w.Close() + return buf, w.FormDataContentType(), err +} + func uploaderTypeByExt(filename string) UploaderFileType { ext := filepath.Ext(filename) switch ext {