From 0b69490757c6d15f9e73a8920e9d462a8f13084d Mon Sep 17 00:00:00 2001
From: Lan Tian <xuyh0120@outlook.com>
Date: Tue, 7 Apr 2020 23:29:48 +0800
Subject: [PATCH] Add telegram bot API

---
 frontend/telegram_bot.go | 117 +++++++++++++++++++++++++++++++++++++++
 frontend/webserver.go    |   1 +
 2 files changed, 118 insertions(+)
 create mode 100644 frontend/telegram_bot.go

diff --git a/frontend/telegram_bot.go b/frontend/telegram_bot.go
new file mode 100644
index 0000000..4667777
--- /dev/null
+++ b/frontend/telegram_bot.go
@@ -0,0 +1,117 @@
+package main
+
+import (
+	"encoding/json"
+	"net/http"
+	"strings"
+)
+
+type tgChat struct {
+	ID int `json:"id"`
+}
+
+type tgMessage struct {
+	MessageID int    `json:"message_id"`
+	Chat      tgChat `json:"chat"`
+	Text      string `json:"text"`
+}
+
+type tgWebhookRequest struct {
+	Message tgMessage `json:"message"`
+}
+
+type tgWebhookResponse struct {
+	Method string `json:"method"`
+	ChatID int    `json:"chat_id"`
+	Text   string `json:"text"`
+}
+
+func telegramIsCommand(message string, command string) bool {
+	b := false
+	b = b || strings.HasPrefix(message, "/"+command+"@")
+	b = b || strings.HasPrefix(message, "/"+command+" ")
+	b = b || message == "/"+command
+	return b
+}
+
+func webHandlerTelegramBot(w http.ResponseWriter, r *http.Request) {
+	// Parse only needed fields of incoming JSON body
+	var err error
+	var request tgWebhookRequest
+	err = json.NewDecoder(r.Body).Decode(&request)
+	if err != nil {
+		println(err.Error())
+		return
+	}
+
+	// Do not respond if not a tg Bot command (starting with /)
+	if request.Message.Text[0] != '/' {
+		return
+	}
+
+	// Select only one server based on webhook URL
+	server := r.URL.Path[len("/telegram/"):]
+	if len(server) == 0 {
+		server = setting.servers[0]
+	}
+
+	// Parse target
+	target := ""
+	if strings.Contains(request.Message.Text, " ") {
+		target = strings.Join(strings.Split(request.Message.Text, " ")[1:], " ")
+	}
+
+	// Execute command
+	commandResult := ""
+
+	// - traceroute
+	if telegramIsCommand(request.Message.Text, "trace") || telegramIsCommand(request.Message.Text, "trace4") {
+		commandResult = batchRequest([]string{server}, "traceroute", target)[0]
+	} else if telegramIsCommand(request.Message.Text, "trace6") {
+		commandResult = batchRequest([]string{server}, "traceroute6", target)[0]
+
+	} else if telegramIsCommand(request.Message.Text, "route") || telegramIsCommand(request.Message.Text, "route4") {
+		commandResult = batchRequest([]string{server}, "bird", "show route for "+target+" primary")[0]
+	} else if telegramIsCommand(request.Message.Text, "route6") {
+		commandResult = batchRequest([]string{server}, "bird6", "show route for "+target+" primary")[0]
+
+	} else if telegramIsCommand(request.Message.Text, "path") || telegramIsCommand(request.Message.Text, "path4") {
+		tempResult := batchRequest([]string{server}, "bird", "show route for "+target+" all primary")[0]
+		for _, s := range strings.Split(tempResult, "\n") {
+			if strings.Contains(s, "BGP.as_path: ") {
+				commandResult = strings.Split(s, "BGP.as_path: ")[1]
+			}
+		}
+	} else if telegramIsCommand(request.Message.Text, "path6") {
+		tempResult := batchRequest([]string{server}, "bird6", "show route for "+target+" all primary")[0]
+		for _, s := range strings.Split(tempResult, "\n") {
+			if strings.Contains(s, "BGP.as_path: ") {
+				commandResult = strings.Split(s, "BGP.as_path: ")[1]
+			}
+		}
+
+	} else if telegramIsCommand(request.Message.Text, "whois") {
+		commandResult = whois(target)
+
+	} else if telegramIsCommand(request.Message.Text, "help") {
+		commandResult = strings.TrimSpace(`
+/[path|path6] <IP>
+/[route|route6] <IP>
+/[trace|trace6] <IP>
+/whois <Target>
+		`)
+	}
+
+	// Create a JSON response
+	w.Header().Add("Content-Type", "application/json")
+	response := &tgWebhookResponse{
+		Method: "sendMessage",
+		ChatID: request.Message.Chat.ID,
+		Text:   commandResult,
+	}
+	err = json.NewEncoder(w).Encode(response)
+	if err != nil {
+		println(err.Error())
+		return
+	}
+}
diff --git a/frontend/webserver.go b/frontend/webserver.go
index 8bd50b5..bffa444 100644
--- a/frontend/webserver.go
+++ b/frontend/webserver.go
@@ -141,5 +141,6 @@ func webServerStart() {
 	http.HandleFunc("/ipv6/traceroute/", webBackendCommunicator("traceroute6", "traceroute"))
 	http.HandleFunc("/whois/", webHandlerWhois)
 	http.HandleFunc("/redir/", webHandlerNavbarFormRedirect)
+	http.HandleFunc("/telegram/", webHandlerTelegramBot)
 	http.ListenAndServe(setting.listen, nil)
 }