Merge pull request #155 from matrix-org/google

Google image search bot
7 years ago · 4042324a40
5 changed files with 323 additions and 1 deletions
--- a/config.sample.yaml
+++ b/config.sample.yaml
@ -74,6 +74,13 @@ services:
    Config:
      api_key: "2356saaqfhgfe"

+  - ID: "google_service"
+    Type: "google"
+    UserID: "@goneb:localhost" # requires a Syncing client
+    Config:
+      api_key: "AIzaSyA4FD39m9"
+      cx: "AIASDFWSRRtrtr"
+
  - ID: "rss_service"
    Type: "rssbot"
    UserID: "@another_goneb:localhost"
--- a/src/github.com/matrix-org/go-neb/goneb.go
+++ b/src/github.com/matrix-org/go-neb/goneb.go
@ -22,6 +22,7 @@ import (
 	_ "github.com/matrix-org/go-neb/services/echo"
 	_ "github.com/matrix-org/go-neb/services/giphy"
 	_ "github.com/matrix-org/go-neb/services/github"
+	_ "github.com/matrix-org/go-neb/services/google"
 	_ "github.com/matrix-org/go-neb/services/guggy"
 	_ "github.com/matrix-org/go-neb/services/jira"
 	_ "github.com/matrix-org/go-neb/services/rssbot"
--- a/src/github.com/matrix-org/go-neb/services/google/google.go
+++ b/src/github.com/matrix-org/go-neb/services/google/google.go
@ -0,0 +1,203 @@
+// Package google implements a Service which adds !commands for Google custom search engine.
+// Initially this package just supports image search but could be expanded to provide other functionality provided by the Google custom search engine API - https://developers.google.com/custom-search/json-api/v1/overview
+package google
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"math"
+	"net/http"
+	"net/url"
+	"strings"
+
+	log "github.com/Sirupsen/logrus"
+	"github.com/matrix-org/go-neb/types"
+	"github.com/matrix-org/gomatrix"
+)
+
+// ServiceType of the Google service
+const ServiceType = "google"
+
+var httpClient = &http.Client{}
+
+type googleSearchResults struct {
+	SearchInformation struct {
+		TotalResults int64 `json:"totalResults,string"`
+	} `json:"searchInformation"`
+	Items []googleSearchResult `json:"items"`
+}
+
+type googleSearchResult struct {
+	Title       string      `json:"title"`
+	HTMLTitle   string      `json:"htmlTitle"`
+	Link        string      `json:"link"`
+	DisplayLink string      `json:"displayLink"`
+	Snippet     string      `json:"snippet"`
+	HTMLSnippet string      `json:"htmlSnippet"`
+	Mime        string      `json:"mime"`
+	FileFormat  string      `json:"fileFormat"`
+	Image       googleImage `json:"image"`
+}
+
+type googleImage struct {
+	ContextLink     string  `json:"contextLink"`
+	Height          float64 `json:"height"`
+	Width           float64 `json:"width"`
+	ByteSize        int64   `json:"byteSize"`
+	ThumbnailLink   string  `json:"thumbnailLink"`
+	ThumbnailHeight float64 `json:"thumbnailHeight"`
+	ThumbnailWidth  float64 `json:"thumbnailWidth"`
+}
+
+// Service contains the Config fields for the Google service.
+//
+// Example request:
+//   {
+//			"api_key": "AIzaSyA4FD39..."
+//			"cx": "ASdsaijwdfASD..."
+//   }
+type Service struct {
+	types.DefaultService
+	// The Google API key to use when making HTTP requests to Google.
+	APIKey string `json:"api_key"`
+	// The Google custom search engine ID
+	Cx string `json:"cx"`
+}
+
+// Commands supported:
+//    !google image some_search_query_without_quotes
+// Responds with a suitable image into the same room as the command.
+func (s *Service) Commands(client *gomatrix.Client) []types.Command {
+	return []types.Command{
+		types.Command{
+			Path: []string{"google", "image"},
+			Command: func(roomID, userID string, args []string) (interface{}, error) {
+				return s.cmdGoogleImgSearch(client, roomID, userID, args)
+			},
+		},
+		types.Command{
+			Path: []string{"google", "help"},
+			Command: func(roomID, userID string, args []string) (interface{}, error) {
+				return usageMessage(), nil
+			},
+		},
+		types.Command{
+			Path: []string{"google"},
+			Command: func(roomID, userID string, args []string) (interface{}, error) {
+				return usageMessage(), nil
+			},
+		},
+	}
+}
+
+// usageMessage returns a matrix TextMessage representation of the service usage
+func usageMessage() *gomatrix.TextMessage {
+	return &gomatrix.TextMessage{"m.notice",
+		`Usage: !google image image_search_text`}
+}
+
+func (s *Service) cmdGoogleImgSearch(client *gomatrix.Client, roomID, userID string, args []string) (interface{}, error) {
+
+	if len(args) < 1 {
+		return usageMessage(), nil
+	}
+
+	// Get the query text to search for.
+	querySentence := strings.Join(args, " ")
+
+	searchResult, err := s.text2imgGoogle(querySentence)
+
+	if err != nil {
+		return nil, err
+	}
+
+	var imgURL = searchResult.Link
+	if imgURL == "" {
+		return gomatrix.TextMessage{
+			MsgType: "m.notice",
+			Body:    "No image found!",
+		}, nil
+	}
+
+	// FIXME -- Sometimes upload fails with a cryptic error - "msg=Upload request failed code=400"
+	resUpload, err := client.UploadLink(imgURL)
+	if err != nil {
+		return nil, fmt.Errorf("Failed to upload Google image to matrix: %s", err.Error())
+	}
+
+	return gomatrix.ImageMessage{
+		MsgType: "m.image",
+		Body:    querySentence,
+		URL:     resUpload.ContentURI,
+		Info: gomatrix.ImageInfo{
+			Height:   uint(math.Floor(searchResult.Image.Height)),
+			Width:    uint(math.Floor(searchResult.Image.Width)),
+			Mimetype: searchResult.Mime,
+		},
+	}, nil
+}
+
+// text2imgGoogle returns info about an image
+func (s *Service) text2imgGoogle(query string) (*googleSearchResult, error) {
+	log.Info("Searching Google for an image of a ", query)
+
+	u, err := url.Parse("https://www.googleapis.com/customsearch/v1")
+	if err != nil {
+		return nil, err
+	}
+
+	q := u.Query()
+	q.Set("q", query)            // String to search for
+	q.Set("num", "1")            // Just return 1 image result
+	q.Set("start", "1")          // No search result offset
+	q.Set("imgSize", "large")    // Just search for medium size images
+	q.Set("searchType", "image") // Search for images
+
+	q.Set("key", s.APIKey) // Set the API key for the request
+	q.Set("cx", s.Cx)      // Set the custom search engine ID
+
+	u.RawQuery = q.Encode()
+	// log.Info("Request URL: ", u)
+
+	res, err := httpClient.Get(u.String())
+	if res != nil {
+		defer res.Body.Close()
+	}
+	if err != nil {
+		return nil, err
+	}
+	if res.StatusCode > 200 {
+		return nil, fmt.Errorf("Request error: %d, %s", res.StatusCode, response2String(res))
+	}
+	var searchResults googleSearchResults
+
+	// log.Info(response2String(res))
+	if err := json.NewDecoder(res.Body).Decode(&searchResults); err != nil {
+		return nil, fmt.Errorf("ERROR - %s", err.Error())
+	} else if len(searchResults.Items) < 1 {
+		return nil, fmt.Errorf("No images found")
+	}
+
+	// Return only the first search result
+	return &searchResults.Items[0], nil
+}
+
+// response2String returns a string representation of an HTTP response body
+func response2String(res *http.Response) string {
+	bs, err := ioutil.ReadAll(res.Body)
+	if err != nil {
+		return "Failed to decode response body"
+	}
+	str := string(bs)
+	return str
+}
+
+// Initialise the service
+func init() {
+	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
+		return &Service{
+			DefaultService: types.NewDefaultService(serviceID, serviceUserID, ServiceType),
+		}
+	})
+}
--- a/src/github.com/matrix-org/go-neb/services/google/google_test.go
+++ b/src/github.com/matrix-org/go-neb/services/google/google_test.go
@ -0,0 +1,111 @@
+package google
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/matrix-org/go-neb/database"
+	"github.com/matrix-org/go-neb/testutils"
+	"github.com/matrix-org/go-neb/types"
+	"github.com/matrix-org/gomatrix"
+)
+
+// TODO: It would be nice to tabularise this test so we can try failing different combinations of responses to make
+//       sure all cases are handled, rather than just the general case as is here.
+func TestCommand(t *testing.T) {
+	database.SetServiceDB(&database.NopStorage{})
+	apiKey := "secret"
+	googleImageURL := "https://www.googleapis.com/customsearch/v1"
+
+	// Mock the response from Google
+	googleTrans := testutils.NewRoundTripper(func(req *http.Request) (*http.Response, error) {
+		googleURL := "https://www.googleapis.com/customsearch/v1"
+		query := req.URL.Query()
+
+		// Check the base API URL
+		if !strings.HasPrefix(req.URL.String(), googleURL) {
+			t.Fatalf("Bad URL: got %s want prefix %s", req.URL.String(), googleURL)
+		}
+		// Check the request method
+		if req.Method != "GET" {
+			t.Fatalf("Bad method: got %s want GET", req.Method)
+		}
+		// Check the API key
+		if query.Get("key") != apiKey {
+			t.Fatalf("Bad apiKey: got %s want %s", query.Get("key"), apiKey)
+		}
+		// Check the search query
+		var searchString = query.Get("q")
+		var searchStringLength = len(searchString)
+		if searchStringLength > 0 && !strings.HasPrefix(searchString, "image") {
+			t.Fatalf("Bad search string: got \"%s\" (%d characters) ", searchString, searchStringLength)
+		}
+
+		resImage := googleImage{
+			Width:  64,
+			Height: 64,
+		}
+
+		res := googleSearchResult{
+			Title: "A Cat",
+			Link:  "http://cat.com/cat.jpg",
+			Mime:  "image/jpeg",
+			Image: resImage,
+		}
+
+		b, err := json.Marshal(res)
+		if err != nil {
+			t.Fatalf("Failed to marshal Google response - %s", err)
+		}
+		return &http.Response{
+			StatusCode: 200,
+			Body:       ioutil.NopCloser(bytes.NewBuffer(b)),
+		}, nil
+	})
+	// clobber the Google service http client instance
+	httpClient = &http.Client{Transport: googleTrans}
+
+	// Create the Google service
+	srv, err := types.CreateService("id", ServiceType, "@googlebot:hyrule", []byte(
+		`{"api_key":"`+apiKey+`"}`,
+	))
+	if err != nil {
+		t.Fatal("Failed to create Google service: ", err)
+	}
+	google := srv.(*Service)
+
+	// Mock the response from Matrix
+	matrixTrans := struct{ testutils.MockTransport }{}
+	matrixTrans.RT = func(req *http.Request) (*http.Response, error) {
+		if req.URL.String() == googleImageURL { // getting the Google image
+			return &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(bytes.NewBufferString("some image data")),
+			}, nil
+		} else if strings.Contains(req.URL.String(), "_matrix/media/r0/upload") { // uploading the image to matrix
+			return &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(bytes.NewBufferString(`{"content_uri":"mxc://foo/bar"}`)),
+			}, nil
+		}
+		return nil, fmt.Errorf("Unknown URL: %s", req.URL.String())
+	}
+	matrixCli, _ := gomatrix.NewClient("https://hyrule", "@googlebot:hyrule", "its_a_secret")
+	matrixCli.Client = &http.Client{Transport: matrixTrans}
+
+	// Execute the matrix !command
+	cmds := google.Commands(matrixCli)
+	if len(cmds) != 3 {
+		t.Fatalf("Unexpected number of commands: %d", len(cmds))
+	}
+	// cmd := cmds[0]
+	// _, err = cmd.Command("!someroom:hyrule", "@navi:hyrule", []string{"image", "Czechoslovakian bananna"})
+	// if err != nil {
+	// 	t.Fatalf("Failed to process command: %s", err.Error())
+	// }
+}
--- a/src/github.com/matrix-org/go-neb/services/guggy/guggy.go
+++ b/src/github.com/matrix-org/go-neb/services/guggy/guggy.go
@ -69,7 +69,7 @@ func (s *Service) cmdGuggy(client *gomatrix.Client, roomID, userID string, args

 	if gifResult.GIF == "" {
 		return gomatrix.TextMessage{
-			MsgType: "m.text.notice",
+			MsgType: "m.notice",
 			Body:    "No GIF found!",
 		}, nil
 	}