sleepyunicorn/main.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/gotd/td/telegram"
	"github.com/gotd/td/telegram/message"
	"github.com/gotd/td/telegram/message/styling"
	"github.com/gotd/td/telegram/uploader"
	"github.com/gotd/td/tg"
	"go.uber.org/zap"
	"golang.org/x/xerrors"
	"gopkg.in/yaml.v2"
)

var (
	config Config
	logger *zap.Logger
)

func init() {
	zapConfig := zap.NewProductionConfig()
	tmpLogger, err := zapConfig.Build()
	if err != nil {
		log.Fatalf("failed to initialise zap: %v", err)
	}
	logger = tmpLogger
	defer logger.Sync()
	configFilename := "config.yaml"
	if len(os.Args) == 2 {
		configFilename = os.Args[1]
	}
	dat, err := os.ReadFile(configFilename)
	if err != nil {
		logger.Fatal("failed to read config.yaml", zap.Error(err))
	}
	err = yaml.Unmarshal(dat, &config)
	if err != nil {
		logger.Fatal("failed to parse config.yaml", zap.Error(err))
	}
}

func ItemInSlice(slice []int, val int) bool {
	for _, i := range slice {
		if i == val {
			return true
		}
	}
	return false
}

func main() {
	defer logger.Sync()
	dispatcher := tg.NewUpdateDispatcher()
	client := telegram.NewClient(config.ApiId, config.ApiHash, telegram.Options{
		Logger:        logger,
		UpdateHandler: dispatcher,
	})
	ctx := context.Background()
	err := client.Run(ctx, func(ctx context.Context) error {
		auth := client.Auth()
		status, err := auth.Status(ctx)
		if err != nil {
			return xerrors.Errorf("failed to check auth status: %v", err)
		}
		if !status.Authorized {
			_, err = auth.Bot(ctx, config.BotToken)
			if err != nil {
				return xerrors.Errorf("failed to authorize: %v", err)
			}
		}
		api := tg.NewClient(client)
		sender := message.NewSender(api)
		dispatcher.OnNewMessage(func(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage) error {
			m, ok := u.Message.(*tg.Message)
			if !ok || m.Out {
				return nil
			}
			switch peer := m.GetPeerID().(type) {
			case *tg.PeerUser:
				if !ItemInSlice(config.AllowedUserIds, peer.UserID) {
					return nil
				}
			default:
				return nil
			}
			split := strings.SplitN(m.Message, "\n", 2)
			if len(split) == 0 {
				return nil
			}
			command := exec.Command("sh", "-c", split[0])
			if len(split) == 2 {
				command.Stdin = strings.NewReader(split[1])
			}
			stdout, err := os.CreateTemp("", "sleepyunicorn*.txt")
			if err != nil {
				return xerrors.Errorf("failed to create temporary file: %v", err)
			}
			defer os.Remove(stdout.Name())
			command.Stdout = stdout
			command.Stderr = stdout
			err = command.Start()
			if err != nil {
				return xerrors.Errorf("failed to start process: %v", err)
			}
			replyUpdate, err := sender.Reply(entities, u).StyledText(ctx, styling.Code(fmt.Sprintf("+++ executing process %d +++", command.Process.Pid)))
			if err != nil {
				return xerrors.Errorf("failed to send executing message: %v", err)
			}
			reply, ok := replyUpdate.(*tg.UpdateShortSentMessage)
			if !ok {
				return xerrors.New("failed to cast to tg.UpdateShortSentMessage")
			}
			err = command.Wait()
			exitCode := 0
			if err != nil {
				if ee, ok := err.(*exec.ExitError); ok {
					exitCode = ee.ExitCode()
				} else {
					return xerrors.Errorf("error occured in command.Run: %v", err)
				}
			}
			exitCodeMessage := fmt.Sprintf("+++ exited with %d +++", exitCode)
			fileSize, err := stdout.Seek(0, 2)
			if err != nil {
				return xerrors.Errorf("failed to seek to end of file: %v", err)
			}
			_, err = stdout.Seek(0, 0)
			if err != nil {
				return xerrors.Errorf("failed to seek to start of file: %v", err)
			}
			reader := bufio.NewReader(stdout)
			if fileSize > 4096-1-int64(len(exitCodeMessage)) {
				upl := uploader.NewUploader(api)
				sender := sender.WithUploader(upl)
				_, err = sender.Answer(entities, u).Edit(reply.ID).StyledText(ctx, styling.Code("+++ uploading output +++"))
				if err != nil {
					return xerrors.Errorf("failed to edit message: %v", err)
				}
				upload, err := upl.FromReader(ctx, "output.txt", reader)
				if err != nil {
					return xerrors.Errorf("failed to upload file: %v", err)
				}
				doc := message.UploadedDocument(upload, styling.Code(exitCodeMessage))
				doc.Filename("output.txt")
				_, err = sender.Reply(entities, u).Media(ctx, doc)
				if err != nil {
					return xerrors.Errorf("failed to send file to chat: %v", err)
				}
				_, err = sender.Answer(entities, u).Revoke().Messages(ctx, reply.ID)
				if err != nil {
					return xerrors.Errorf("failed to delete message: %v", err)
				}
				return nil
			}
			var builder strings.Builder
			_, err = reader.WriteTo(&builder)
			if err != nil {
				return xerrors.Errorf("error occured in bufio.Reader.WriteTo: %v", err)
			}
			_, err = sender.Answer(entities, u).Edit(reply.ID).StyledText(ctx, styling.Code(exitCodeMessage), styling.Plain("\n"), styling.Pre(builder.String()))
			if err != nil {
				return xerrors.Errorf("failed to edit message: %v", err)
			}
			return nil
		})
		return telegram.RunUntilCanceled(ctx, client)
	})
	if err != nil {
		logger.Fatal("error occured in client.Run", zap.Error(err))
	}
}