[tgbot] Ability to select telegram user for client from bot

2 files changed, 153 insertions, 3 deletions

diff --git a/web/service/inbound.go b/web/service/inbound.go
index 7e838b5f..57df0b9e 100644
--- a/web/service/inbound.go
+++ b/web/service/inbound.go
@@ -673,6 +673,21 @@ func (s *InboundService) DelClientIPs(tx *gorm.DB, email string) error {
 	return tx.Where("client_email = ?", email).Delete(model.InboundClientIps{}).Error
 }
 
+func (s *InboundService) GetClientInboundByTrafficID(trafficId int) (traffic *xray.ClientTraffic, inbound *model.Inbound, err error) {
+	db := database.GetDB()
+	var traffics []*xray.ClientTraffic
+	err = db.Model(xray.ClientTraffic{}).Where("id = ?", trafficId).Find(&traffics).Error
+	if err != nil {
+		logger.Warning(err)
+		return nil, nil, err
+	}
+	if len(traffics) > 0 {
+		inbound, err = s.GetInbound(traffics[0].InboundId)
+		return traffics[0], inbound, err
+	}
+	return nil, nil, nil
+}
+
 func (s *InboundService) GetClientInboundByEmail(email string) (traffic *xray.ClientTraffic, inbound *model.Inbound, err error) {
 	db := database.GetDB()
 	var traffics []*xray.ClientTraffic
@@ -688,6 +703,85 @@ func (s *InboundService) GetClientInboundByEmail(email string) (traffic *xray.Cl
 	return nil, nil, nil
 }
 
+func (s *InboundService) GetClientByEmail(clientEmail string) (*xray.ClientTraffic, *model.Client, error) {
+	traffic, inbound, err := s.GetClientInboundByEmail(clientEmail)
+	if err != nil {
+		return nil, nil, err
+	}
+	if inbound == nil {
+		return nil, nil, common.NewError("Inbound Not Found For Email:", clientEmail)
+	}
+
+	clients, err := s.getClients(inbound)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	for _, client := range clients {
+		if client.Email == clientEmail {
+			return traffic, &client, nil
+		}
+	}
+
+	return nil, nil, common.NewError("Client Not Found In Inbound For Email:", clientEmail)
+}
+
+func (s *InboundService) SetClientTelegramUserID(trafficId int, tgId string) error {
+	traffic, inbound, err := s.GetClientInboundByTrafficID(trafficId)
+	if err != nil {
+		return err
+	}
+	if inbound == nil {
+		return common.NewError("Inbound Not Found For Traffic ID:", trafficId)
+	}
+
+	clientEmail := traffic.Email
+
+	oldClients, err := s.getClients(inbound)
+	if err != nil {
+		return err
+	}
+
+	clientId := ""
+
+	for _, oldClient := range oldClients {
+		if oldClient.Email == clientEmail {
+			if inbound.Protocol == "trojan" {
+				clientId = oldClient.Password
+			} else {
+				clientId = oldClient.ID
+			}
+			break
+		}
+	}
+
+	if len(clientId) == 0 {
+		return common.NewError("Client Not Found For Email:", clientEmail)
+	}
+
+	var settings map[string]interface{}
+	err = json.Unmarshal([]byte(inbound.Settings), &settings)
+	if err != nil {
+		return err
+	}
+	clients := settings["clients"].([]interface{})
+	var newClients []interface{}
+	for client_index := range clients {
+		c := clients[client_index].(map[string]interface{})
+		if c["email"] == clientEmail {
+			c["tgId"] = tgId
+			newClients = append(newClients, interface{}(c))
+		}
+	}
+	settings["clients"] = newClients
+	modifiedSettings, err := json.MarshalIndent(settings, "", "  ")
+	if err != nil {
+		return err
+	}
+	inbound.Settings = string(modifiedSettings)
+	return s.UpdateInboundClient(inbound, clientId)
+}
+
 func (s *InboundService) ToggleClientEnableByEmail(clientEmail string) (bool, error) {
 	_, inbound, err := s.GetClientInboundByEmail(clientEmail)
 	if err != nil {
diff --git a/web/service/tgbot.go b/web/service/tgbot.go
index e5f35cbf..3168839b 100644
--- a/web/service/tgbot.go
+++ b/web/service/tgbot.go
@@ -102,6 +102,10 @@ func (t *Tgbot) OnReceive() {
 	botHandler, _ = th.NewBotHandler(bot, updates)
 
 	botHandler.HandleMessage(func(bot *telego.Bot, message telego.Message) {
+		t.SendMsgToTgbot(message.Chat.ID, "Custom Keyboard Closed!", tu.ReplyKeyboardRemove())
+	}, th.TextEqual("❌ Close Keyboard"))
+
+	botHandler.HandleMessage(func(bot *telego.Bot, message telego.Message) {
 		t.answerCommand(&message, message.Chat.ID, checkAdmin(message.From.ID))
 	}, th.AnyCommand())
 
@@ -109,6 +113,23 @@ func (t *Tgbot) OnReceive() {
 		t.asnwerCallback(&query, checkAdmin(query.From.ID))
 	}, th.AnyCallbackQueryWithMessage())
 
+	botHandler.HandleMessage(func(bot *telego.Bot, message telego.Message) {
+		if message.UserShared != nil {
+			if checkAdmin(message.From.ID) {
+				err := t.inboundService.SetClientTelegramUserID(message.UserShared.RequestID, strconv.FormatInt(message.UserShared.UserID, 10))
+				var output string
+				if err != nil {
+					output = "❌ Error in user selection!"
+				} else {
+					output = "✅ Telegram User saved."
+				}
+				t.SendMsgToTgbot(message.Chat.ID, output, tu.ReplyKeyboardRemove())
+			} else {
+				t.SendMsgToTgbot(message.Chat.ID, "No result!", tu.ReplyKeyboardRemove())
+			}
+		}
+	}, th.AnyMessage())
+
 	botHandler.Start()
 }
 
@@ -301,6 +322,9 @@ func (t *Tgbot) asnwerCallback(callbackQuery *telego.CallbackQuery, isAdmin bool
 			case "ip_log":
 				t.sendCallbackAnswerTgBot(callbackQuery.ID, fmt.Sprintf("✅ %s : Get IP Log.", email))
 				t.searchClientIps(chatId, email)
+			case "tg_user":
+				t.sendCallbackAnswerTgBot(callbackQuery.ID, fmt.Sprintf("✅ %s : Get Telegram User Info.", email))
+				t.clientTelegramUserInfo(chatId, email)
 			case "toggle_enable":
 				enabled, err := t.inboundService.ToggleClientEnableByEmail(email)
 				if err == nil {
@@ -386,7 +410,7 @@ func (t *Tgbot) SendAnswer(chatId int64, msg string, isAdmin bool) {
 	}
 }
 
-func (t *Tgbot) SendMsgToTgbot(chatId int64, msg string, inlineKeyboard ...*telego.InlineKeyboardMarkup) {
+func (t *Tgbot) SendMsgToTgbot(chatId int64, msg string, replyMarkup ...telego.ReplyMarkup) {
 	if !isRunning {
 		return
 	}
@@ -413,8 +437,8 @@ func (t *Tgbot) SendMsgToTgbot(chatId int64, msg string, inlineKeyboard ...*tele
 			Text:      message,
 			ParseMode: "HTML",
 		}
-		if len(inlineKeyboard) > 0 {
-			params.ReplyMarkup = inlineKeyboard[0]
+		if len(replyMarkup) > 0 {
+			params.ReplyMarkup = replyMarkup[0]
 		}
 		_, err := bot.SendMessage(&params)
 		if err != nil {
@@ -613,6 +637,35 @@ func (t *Tgbot) searchClientIps(chatId int64, email string, messageID ...int) {
 	}
 }
 
+func (t *Tgbot) clientTelegramUserInfo(chatId int64, email string) {
+	traffic, client, err := t.inboundService.GetClientByEmail(email)
+	if err != nil {
+		logger.Warning(err)
+		msg := "❌ Something went wrong!"
+		t.SendMsgToTgbot(chatId, msg)
+		return
+	}
+	if client == nil {
+		msg := "No result!"
+		t.SendMsgToTgbot(chatId, msg)
+		return
+	}
+	output := fmt.Sprintf("📧 Email: %s\r\n👤 Telegram User: %s\r\n", email, client.TgID)
+	requestUser := telego.KeyboardButtonRequestUser{
+		RequestID: int32(traffic.Id),
+		UserIsBot: false,
+	}
+	keyboard := tu.Keyboard(
+		tu.KeyboardRow(
+			tu.KeyboardButton("👤 Select Telegram User").WithRequestUser(&requestUser),
+		),
+		tu.KeyboardRow(
+			tu.KeyboardButton("❌ Close Keyboard"),
+		),
+	).WithIsPersistent()
+	t.SendMsgToTgbot(chatId, output, keyboard)
+}
+
 func (t *Tgbot) searchClient(chatId int64, email string, messageID ...int) {
 	traffic, err := t.inboundService.GetClientTrafficByEmail(email)
 	if err != nil {
@@ -658,6 +711,9 @@ func (t *Tgbot) searchClient(chatId int64, email string, messageID ...int) {
 			tu.InlineKeyboardButton("🔢 IP Limit").WithCallbackData("ip_limit "+email),
 		),
 		tu.InlineKeyboardRow(
+			tu.InlineKeyboardButton("👤 Set Telegram User").WithCallbackData("tg_user "+email),
+		),
+		tu.InlineKeyboardRow(
 			tu.InlineKeyboardButton("🔘 Enable / Disable").WithCallbackData("toggle_enable "+email),
 		),
 	)