muskrat/bot/handlers/code.go

package handlers

import (
	"bytes"
	"context"
	"errors"
	"gitea.russia9.dev/russia9/muskrat/storage"
	"gitea.russia9.dev/russia9/muskrat/types"
	"gitea.russia9.dev/russia9/muskrat/userbot"
	"github.com/gotd/td/session"
	"github.com/gotd/td/telegram"
	"github.com/gotd/td/telegram/auth"
	"gopkg.in/telebot.v3"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"
)

func Code(ctx telebot.Context, metadata []string) error {
	if !regexp.MustCompile("^\\d \\d \\d \\d \\d$").MatchString(ctx.Text()) {
		_ = storage.ClearState(context.Background(), ctx.Sender().ID)
		return ctx.Send("Неправильный формат кода. <b>Научись читать пж</b>")
	}

	sess := storage.Registration[ctx.Sender().ID]
	if len(sess) == 0 {
		_ = storage.ClearState(context.Background(), ctx.Sender().ID)
		return List(ctx)
	}

	s := session.StorageMemory{}
	err := s.StoreSession(context.Background(), sess)
	if err != nil {
		_ = storage.ClearState(context.Background(), ctx.Sender().ID)
		return err
	}

	appID, err := strconv.Atoi(os.Getenv("TELEGRAM_APP_ID"))
	if err != nil {
		return err
	}
	client := telegram.NewClient(appID, os.Getenv("TELEGRAM_APP_HASH"), telegram.Options{
		SessionStorage: &s,
	})

	err = client.Run(context.Background(), func(c context.Context) error {
		_, err = client.Auth().SignIn(c, metadata[0], strings.ReplaceAll(ctx.Text(), " ", ""), metadata[1])
		if errors.Is(err, auth.ErrPasswordAuthNeeded) {
			b := new(bytes.Buffer)
			err = s.Dump(b)
			if err != nil {
				_ = storage.ClearState(context.Background(), ctx.Sender().ID)
				return err
			}
			if len(storage.Registration) == 0 {
				storage.Registration = make(map[int64][]byte)
			}
			storage.Registration[ctx.Sender().ID] = b.Bytes()
			storage.SetState(context.Background(), ctx.Sender().ID, "password")
			return ctx.Send("Введите <b>пароль 2fa</b>:")
		} else if err != nil {
			return err
		}

		self, err := client.Self(context.Background())
		if err != nil {
			return err
		}

		// Success
		account := &types.Account{
			ID:      self.ID,
			User:    ctx.Sender().ID,
			Session: storage.Registration[ctx.Sender().ID],
		}
		err = storage.DB.Create(account).Error
		if err != nil {
			return ctx.Send("Аккаунт с таким ID уже есть в боте")
		}

		if len(userbot.Userbots) == 0 {
			userbot.Userbots = make(map[int64]*userbot.Userbot)
		}
		if userbot.Userbots[account.ID] != nil {
			userbot.Userbots[account.ID].Stop()
			delete(userbot.Userbots, account.ID)
		}
		ubot := &userbot.Userbot{
			Account: account,
		}
		go ubot.Start(c)
		userbot.Userbots[account.ID] = ubot

		time.Sleep(time.Second * 2)

		delete(storage.Registration, ctx.Sender().ID)
		_ = storage.ClearState(context.Background(), ctx.Sender().ID)

		return ctx.Send("Успешный логин" + self.FirstName + " " + self.LastName)
	})

	return err
}