mirror of https://github.com/matrix-org/go-neb.git
Richard Lewis
7 years ago
4 changed files with 335 additions and 0 deletions
-
6config.sample.yaml
-
1src/github.com/matrix-org/go-neb/goneb.go
-
221src/github.com/matrix-org/go-neb/services/google/google.go
-
107src/github.com/matrix-org/go-neb/services/google/google_test.go
@ -0,0 +1,221 @@ |
|||
// 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{} |
|||
|
|||
// Unsused -- leaving this in place for the time being to show structure of the request
|
|||
// type googleQuery struct {
|
|||
// // Query search text
|
|||
// Query string `json:"q"`
|
|||
// // Number of search results
|
|||
// Num int `json:"num"`
|
|||
// // Search result offset
|
|||
// Start int `json:"start"`
|
|||
// // Size of images to serch for (usually set to "medium")
|
|||
// ImgSize string `json:"imgSize"`
|
|||
// // Type of search - Currently always set to "image"
|
|||
// SearchType string `json:"searchType"`
|
|||
// // Type of image file to retur64 `json:"totalResults"`
|
|||
// FileType string `json:"fileType"`
|
|||
// // API key
|
|||
// Key string `json:"key"`
|
|||
// // Custom serch engine ID
|
|||
// Cx string `json:"cx"`
|
|||
// }
|
|||
|
|||
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.
|
|||
// TODO - move the google custom search engine ID in here!
|
|||
//
|
|||
// Example request:
|
|||
// {
|
|||
// "api_key": "AIzaSyA4FD39m9pN-hiYf2NRU9x9cOv5tekRDvM"
|
|||
// }
|
|||
type Service struct { |
|||
types.DefaultService |
|||
// The Google API key to use when making HTTP requests to Google.
|
|||
APIKey string `json:"api_key"` |
|||
} |
|||
|
|||
// Commands supported:
|
|||
// !google 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"}, |
|||
Command: func(roomID, userID string, args []string) (interface{}, error) { |
|||
return s.cmdGoogle(client, roomID, userID, args) |
|||
}, |
|||
}, |
|||
} |
|||
} |
|||
|
|||
// 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) cmdGoogle(client *gomatrix.Client, roomID, userID string, args []string) (interface{}, error) { |
|||
|
|||
if len(args) < 2 || args[0] != "image" { |
|||
return usageMessage(), nil |
|||
} |
|||
// Drop the search type (should currently always be "image")
|
|||
args = args[1:] |
|||
|
|||
// only 1 arg which is the 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.text.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()) |
|||
} |
|||
|
|||
img := searchResult.Image |
|||
return gomatrix.ImageMessage{ |
|||
MsgType: "m.image", |
|||
Body: querySentence, |
|||
URL: resUpload.ContentURI, |
|||
Info: gomatrix.ImageInfo{ |
|||
Height: uint(math.Floor(img.Height)), |
|||
Width: uint(math.Floor(img.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", "medium") // Just search for medium size images
|
|||
q.Set("searchType", "image") // Search for images
|
|||
// q.set("fileType, "") // Any file format
|
|||
|
|||
var key = s.APIKey |
|||
if key == "" { |
|||
key = "AIzaSyA4FD39m9pN-hiYf2NRU9x9cOv5tekRDvM" // FIXME -- Should be instantiated from service config
|
|||
} |
|||
q.Set("key", key) // Set the API key for the request
|
|||
q.Set("cx", "003141582324323361145:f5zyrk9_8_m") // Set the custom search engine ID
|
|||
|
|||
u.RawQuery = q.Encode() |
|||
// log.Info("Request URL: ", u)
|
|||
|
|||
res, err := http.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 || len(searchResults.Items) < 1 { |
|||
// Google return a JSON object which has { items: [] } if there are 0 results.
|
|||
// This fails to be deserialised by Go.
|
|||
|
|||
// TODO -- Find out how to just return an error string (with no formatting)
|
|||
// return nil, errors.New("No images found")
|
|||
// return nil, fmt.Errorf("No results - %s", err)
|
|||
return nil, fmt.Errorf("No images found%s", "") |
|||
} |
|||
|
|||
// 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) (responseText 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), |
|||
} |
|||
}) |
|||
} |
|||
@ -0,0 +1,107 @@ |
|||
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 custom search engine ID
|
|||
// Check the search query
|
|||
|
|||
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) != 1 { |
|||
t.Fatalf("Unexpected number of commands: %d", len(cmds)) |
|||
} |
|||
cmd := cmds[0] |
|||
_, err = cmd.Command("!someroom:hyrule", "@navi:hyrule", []string{"Czechoslovakian", "bananna"}) |
|||
if err != nil { |
|||
t.Fatalf("Failed to process command: %s", err.Error()) |
|||
} |
|||
} |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue