Add safe-search option to google service

.gitignore

@@ -1,3 +1,4 @@
+config.yaml
 godoc
 go-neb.db
 .*.swp

config.sample.yaml

@@ -81,6 +81,7 @@ services:
     Config:
       api_key: "AIzaSyA4FD39m9"
       cx: "AIASDFWSRRtrtr"
+      safe_search: false
 
   - ID: "imgur_service"
     Type: "imgur"

src/github.com/matrix-org/go-neb/services/google/google.go

@@ -3,8 +3,10 @@
 package google
 
 import (
+	"database/sql"
 	"encoding/json"
 	"fmt"
+	"github.com/matrix-org/go-neb/database"
 	"io/ioutil"
 	"math"
 	"net/http"
@@ -57,12 +59,25 @@ type googleImage struct {
 //			"api_key": "AIzaSyA4FD39..."
 //			"cx": "ASdsaijwdfASD..."
 //   }
+//
+// You can set a room-specific safe-search enabled flag for a Matrix room by sending a `m.room.bot.options` state event
+// which has the following `content`:
+//
+//  {
+//    "google": {
+//      "safe_search": false
+//    }
+//  }
+//
+// This will override the global Safesearch setting on a per-room basis.
 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"`
+	// Whether or not to enable Google safe-search
+	Safesearch bool `json:"safe_search"`
 }
 
 // Commands supported:
@@ -106,7 +121,7 @@ func (s *Service) cmdGoogleImgSearch(client *gomatrix.Client, roomID, userID str
 	// Get the query text to search for.
 	querySentence := strings.Join(args, " ")
 
-	searchResult, err := s.text2imgGoogle(querySentence)
+	searchResult, err := s.text2imgGoogle(roomID, querySentence)
 
 	if err != nil {
 		return nil, err
@@ -138,8 +153,38 @@ func (s *Service) cmdGoogleImgSearch(client *gomatrix.Client, roomID, userID str
 	}, nil
 }
 
+// safeSearchEnabled returns whether safe search is enabled for the given room, or the service global value s.Safesearch
+func (s *Service) safeSearchEnabled(roomID string) bool {
+	logger := log.WithFields(log.Fields{
+		"room_id":     roomID,
+		"bot_user_id": s.ServiceUserID(),
+	})
+	opts, err := database.GetServiceDB().LoadBotOptions(s.ServiceUserID(), roomID)
+	if err != nil {
+		if err != sql.ErrNoRows {
+			logger.WithError(err).Error("Failed to load bot options")
+		}
+		return s.Safesearch
+	}
+	// Expect opts to look like:
+	// { google: { safe_search: BOOLEAN } }
+	googleOpts, ok := opts.Options["google"].(map[string]interface{})
+	if !ok {
+		logger.WithField("options", opts.Options).Error("Failed to cast bot options as google options")
+		return s.Safesearch
+	}
+	safeSearch, ok := googleOpts["safe_search"].(bool)
+	if !ok {
+		logger.WithField("default_repo", googleOpts["safe_search"]).Error(
+			"Failed to cast room-specific safe-search as a bool",
+		)
+		return s.Safesearch
+	}
+	return safeSearch
+}
+
 // text2imgGoogle returns info about an image
-func (s *Service) text2imgGoogle(query string) (*googleSearchResult, error) {
+func (s *Service) text2imgGoogle(roomID, query string) (*googleSearchResult, error) {
 	log.Info("Searching Google for an image of a ", query)
 
 	u, err := url.Parse("https://www.googleapis.com/customsearch/v1")
@@ -157,6 +202,12 @@ func (s *Service) text2imgGoogle(query string) (*googleSearchResult, error) {
 	q.Set("key", s.APIKey) // Set the API key for the request
 	q.Set("cx", s.Cx)      // Set the custom search engine ID
 
+	if s.safeSearchEnabled(roomID) {
+		q.Set("safe", "active")
+	} else {
+		q.Set("safe", "off")
+	}
+
 	u.RawQuery = q.Encode()
 	// log.Info("Request URL: ", u)
 

src/github.com/matrix-org/go-neb/services/google/google_test.go

@@ -15,15 +15,8 @@ import (
 	"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 := "http://cat.com/cat.jpg"
-
-	// Mock the response from Google
-	googleTrans := testutils.NewRoundTripper(func(req *http.Request) (*http.Response, error) {
+func mockGoogle(t *testing.T, apiKey, googleImageURL string, safeSearch bool) http.RoundTripper {
+	return testutils.NewRoundTripper(func(req *http.Request) (*http.Response, error) {
 		googleURL := "https://www.googleapis.com/customsearch/v1"
 		query := req.URL.Query()
 
@@ -39,6 +32,15 @@ func TestCommand(t *testing.T) {
 		if query.Get("key") != apiKey {
 			t.Fatalf("Bad apiKey: got %s want %s", query.Get("key"), apiKey)
 		}
+		// Check safe-search
+		safe := "off"
+		if safeSearch {
+			safe = "active"
+		}
+		if query.Get("safe") != safe {
+			t.Fatalf("Bad safe: got %s want %s", query.Get("safe"), safe)
+		}
+
 		// Check the search query
 		var searchString = query.Get("q")
 		var searchStringLength = len(searchString)
@@ -73,6 +75,17 @@ func TestCommand(t *testing.T) {
 			Body:       ioutil.NopCloser(bytes.NewBuffer(b)),
 		}, nil
 	})
+}
+
+// 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 := "http://cat.com/cat.jpg"
+
+	// Mock the response from Google
+	googleTrans := mockGoogle(t, apiKey, googleImageURL, false)
 	// clobber the Google service http client instance
 	httpClient = &http.Client{Transport: googleTrans}
 
@@ -115,3 +128,53 @@ func TestCommand(t *testing.T) {
 		t.Fatalf("Failed to process command: %s", err.Error())
 	}
 }
+
+func TestSafesearch(t *testing.T) {
+	database.SetServiceDB(&database.NopStorage{})
+	apiKey := "secret"
+	googleImageURL := "http://cat.com/cat.jpg"
+
+	// Mock the response from Google
+	googleTrans := mockGoogle(t, apiKey, googleImageURL, true)
+	// 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+`", "safe_search": true}`,
+	))
+	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())
+	}
+}