package main import ( "bytes" "encoding/json" "net/http" "net/http/httptest" "net/url" "strings" "testing" "github.com/jarcoal/httpmock" "github.com/magiconair/properties/assert" ) func doTestTelegramIsCommand(t *testing.T, message string, command string, expected bool) { result := telegramIsCommand(message, command) assert.Equal(t, result, expected) } func mockTelegramCall(t *testing.T, msg string, raw bool) string { return mockTelegramEndpointCall(t, "/telegram/", msg, raw) } func mockTelegramEndpointCall(t *testing.T, endpoint string, msg string, raw bool) string { request := tgWebhookRequest{ Message: tgMessage{ MessageID: 123, Chat: tgChat{ ID: 456, }, Text: msg, }, } requestJson, err := json.Marshal(request) if err != nil { t.Fatal(err) } requestBody := bytes.NewReader(requestJson) r := httptest.NewRequest(http.MethodGet, endpoint, requestBody) w := httptest.NewRecorder() webHandlerTelegramBot(w, r) if raw { return w.Body.String() } else { var response tgWebhookResponse if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { t.Error(err) } assert.Equal(t, response.ChatID, request.Message.Chat.ID) assert.Equal(t, response.ReplyToMessageID, request.Message.MessageID) return response.Text } } func TestTelegramIsCommand(t *testing.T) { setting.telegramBotName = "test_bot" // Recognize command doTestTelegramIsCommand(t, "/trace", "trace", true) doTestTelegramIsCommand(t, "/trace", "trace1234", false) doTestTelegramIsCommand(t, "/trace", "tra", false) doTestTelegramIsCommand(t, "/trace", "abcdefg", false) // Recognize command with parameters doTestTelegramIsCommand(t, "/trace google.com", "trace", true) doTestTelegramIsCommand(t, "/trace google.com", "trace1234", false) doTestTelegramIsCommand(t, "/trace google.com", "tra", false) doTestTelegramIsCommand(t, "/trace google.com", "abcdefg", false) // Recognize command with bot name doTestTelegramIsCommand(t, "/trace@test_bot", "trace", true) doTestTelegramIsCommand(t, "/trace@test_bot", "trace1234", false) doTestTelegramIsCommand(t, "/trace@test_bot", "tra", false) doTestTelegramIsCommand(t, "/trace@test_bot", "abcdefg", false) doTestTelegramIsCommand(t, "/trace@test_bot_123", "trace", false) doTestTelegramIsCommand(t, "/trace@test_", "trace", false) // Recognize command with bot name and parameters doTestTelegramIsCommand(t, "/trace@test_bot google.com", "trace", true) doTestTelegramIsCommand(t, "/trace@test_bot google.com", "trace1234", false) doTestTelegramIsCommand(t, "/trace@test_bot google.com", "tra", false) doTestTelegramIsCommand(t, "/trace@test_bot google.com", "abcdefg", false) doTestTelegramIsCommand(t, "/trace@test_bot_123 google.com", "trace", false) doTestTelegramIsCommand(t, "/trace@test google.com", "trace", false) } func TestTelegramBatchRequestFormatSingleServer(t *testing.T) { httpmock.Activate() defer httpmock.DeactivateAndReset() httpResponse := httpmock.NewStringResponder(200, "Mock") httpmock.RegisterResponder("GET", "http://alpha:8000/mock?q=cmd", httpResponse) setting.servers = []string{"alpha"} setting.domain = "" setting.proxyPort = 8000 result := telegramBatchRequestFormat(setting.servers, "mock", "cmd", telegramDefaultPostProcess) expected := "Mock\n\n" assert.Equal(t, result, expected) } func TestTelegramBatchRequestFormatMultipleServers(t *testing.T) { httpmock.Activate() defer httpmock.DeactivateAndReset() httpResponse := httpmock.NewStringResponder(200, "Mock") httpmock.RegisterResponder("GET", "http://alpha:8000/mock?q=cmd", httpResponse) httpmock.RegisterResponder("GET", "http://beta:8000/mock?q=cmd", httpResponse) httpmock.RegisterResponder("GET", "http://gamma:8000/mock?q=cmd", httpResponse) setting.servers = []string{ "alpha", "beta", "gamma", } setting.domain = "" setting.proxyPort = 8000 result := telegramBatchRequestFormat(setting.servers, "mock", "cmd", telegramDefaultPostProcess) expected := "alpha\nMock\n\nbeta\nMock\n\ngamma\nMock\n\n" assert.Equal(t, result, expected) } func TestWebHandlerTelegramBotBadJSON(t *testing.T) { requestBody := strings.NewReader("{bad json}") r := httptest.NewRequest(http.MethodGet, "/telegram/", requestBody) w := httptest.NewRecorder() webHandlerTelegramBot(w, r) response := w.Body.String() assert.Equal(t, response, "") } func TestWebHandlerTelegramBotTrace(t *testing.T) { httpmock.Activate() defer httpmock.DeactivateAndReset() httpResponse := httpmock.NewStringResponder(200, "Mock Response") httpmock.RegisterResponder("GET", "http://alpha:8000/traceroute?q=1.1.1.1", httpResponse) setting.servers = []string{"alpha"} setting.domain = "" setting.proxyPort = 8000 response := mockTelegramCall(t, "/trace 1.1.1.1", false) assert.Equal(t, response, "```\nMock Response\n```") } func TestWebHandlerTelegramBotTraceWithServerList(t *testing.T) { httpmock.Activate() defer httpmock.DeactivateAndReset() httpResponse := httpmock.NewStringResponder(200, "Mock Response") httpmock.RegisterResponder("GET", "http://alpha:8000/traceroute?q=1.1.1.1", httpResponse) setting.servers = []string{"alpha"} setting.domain = "" setting.proxyPort = 8000 response := mockTelegramEndpointCall(t, "/telegram/alpha", "/trace 1.1.1.1", false) assert.Equal(t, response, "```\nMock Response\n```") } func TestWebHandlerTelegramBotRoute(t *testing.T) { httpmock.Activate() defer httpmock.DeactivateAndReset() httpResponse := httpmock.NewStringResponder(200, "Mock Response") httpmock.RegisterResponder("GET", "http://alpha:8000/bird?q="+url.QueryEscape("show route for 1.1.1.1 primary"), httpResponse) setting.servers = []string{"alpha"} setting.domain = "" setting.proxyPort = 8000 response := mockTelegramCall(t, "/route 1.1.1.1", false) assert.Equal(t, response, "```\nMock Response\n```") } func TestWebHandlerTelegramBotPath(t *testing.T) { httpmock.Activate() defer httpmock.DeactivateAndReset() httpResponse := httpmock.NewStringResponder(200, ` BGP.as_path: 123 456 `) httpmock.RegisterResponder("GET", "http://alpha:8000/bird?q="+url.QueryEscape("show route for 1.1.1.1 all primary"), httpResponse) setting.servers = []string{"alpha"} setting.domain = "" setting.proxyPort = 8000 response := mockTelegramCall(t, "/path 1.1.1.1", false) assert.Equal(t, response, "```\n123 456\n```") } func TestWebHandlerTelegramBotPathMissing(t *testing.T) { httpmock.Activate() defer httpmock.DeactivateAndReset() httpResponse := httpmock.NewStringResponder(200, "No path in this response") httpmock.RegisterResponder("GET", "http://alpha:8000/bird?q="+url.QueryEscape("show route for 1.1.1.1 all primary"), httpResponse) setting.servers = []string{"alpha"} setting.domain = "" setting.proxyPort = 8000 response := mockTelegramCall(t, "/path 1.1.1.1", false) assert.Equal(t, response, "```\nempty result\n```") } func TestWebHandlerTelegramBotWhois(t *testing.T) { server := WhoisServer{ t: t, expectedQuery: "AS6939", response: AS6939Response, } server.Listen() go server.Run() defer server.Close() setting.netSpecificMode = "" setting.whoisServer = server.server.Addr().String() response := mockTelegramCall(t, "/whois AS6939", false) assert.Equal(t, response, "```"+server.response+"```") } func TestWebHandlerTelegramBotWhoisDN42Mode(t *testing.T) { server := WhoisServer{ t: t, expectedQuery: "AS4242422547", response: ` Query for AS4242422547 `, } server.Listen() go server.Run() defer server.Close() setting.netSpecificMode = "dn42" setting.whoisServer = server.server.Addr().String() response := mockTelegramCall(t, "/whois 2547", false) assert.Equal(t, response, "```"+server.response+"```") } func TestWebHandlerTelegramBotWhoisDN42ModeFullASN(t *testing.T) { server := WhoisServer{ t: t, expectedQuery: "AS4242422547", response: ` Query for AS4242422547 `, } server.Listen() go server.Run() defer server.Close() setting.netSpecificMode = "dn42" setting.whoisServer = server.server.Addr().String() response := mockTelegramCall(t, "/whois 4242422547", false) assert.Equal(t, response, "```"+server.response+"```") } func TestWebHandlerTelegramBotWhoisShortenMode(t *testing.T) { server := WhoisServer{ t: t, expectedQuery: "AS6939", response: ` Information line that will be removed # Comment that will be removed Name: Redacted for privacy Descr: This is a vvvvvvvvvvvvvvvvvvvvvvveeeeeeeeeeeeeeeeeeeerrrrrrrrrrrrrrrrrrrrrrrryyyyyyyyyyyyyyyyyyy long line that will be skipped. Looooooooooooooooooooooong key: this line will be skipped. Preserved1: this line isn't removed. Preserved2: this line isn't removed. Preserved3: this line isn't removed. Preserved4: this line isn't removed. Preserved5: this line isn't removed. `, } expectedResult := `Preserved1: this line isn't removed. Preserved2: this line isn't removed. Preserved3: this line isn't removed. Preserved4: this line isn't removed. Preserved5: this line isn't removed. 3 line(s) skipped.` server.Listen() go server.Run() defer server.Close() setting.netSpecificMode = "shorten" setting.whoisServer = server.server.Addr().String() response := mockTelegramCall(t, "/whois AS6939", false) assert.Equal(t, response, "```\n"+expectedResult+"\n```") } func TestWebHandlerTelegramBotHelp(t *testing.T) { response := mockTelegramCall(t, "/help", false) if !strings.Contains(response, "/trace") { t.Error("Did not get help message") } } func TestWebHandlerTelegramBotUnknownCommand(t *testing.T) { response := mockTelegramCall(t, "/nonexistent", true) assert.Equal(t, response, "") } func TestWebHandlerTelegramBotNotCommand(t *testing.T) { response := mockTelegramCall(t, "random chat message", true) assert.Equal(t, response, "") } func TestWebHandlerTelegramBotEmptyResponse(t *testing.T) { server := WhoisServer{ t: t, expectedQuery: "AS6939", response: "", } server.Listen() go server.Run() defer server.Close() setting.netSpecificMode = "" setting.whoisServer = server.server.Addr().String() response := mockTelegramCall(t, "/whois AS6939", false) assert.Equal(t, response, "```\nempty result\n```") } func TestWebHandlerTelegramBotTruncateLongResponse(t *testing.T) { server := WhoisServer{ t: t, expectedQuery: "AS6939", response: strings.Repeat("A", 65536), } server.Listen() go server.Run() defer server.Close() setting.netSpecificMode = "" setting.whoisServer = server.server.Addr().String() response := mockTelegramCall(t, "/whois AS6939", false) assert.Equal(t, response, "```\n"+strings.Repeat("A", 4096)+"\n```") }