From 54fe866d7131fe3eb83846c3a5b03c861531bed1 Mon Sep 17 00:00:00 2001 From: Jordan Wright Date: Sun, 23 Aug 2015 20:42:47 -0500 Subject: [PATCH] Added /api/import/site functionality --- controllers/api.go | 15 +++++++-- controllers/api_test.go | 70 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 2 deletions(-) create mode 100644 controllers/api_test.go diff --git a/controllers/api.go b/controllers/api.go index a4fdcc8f..d281d032 100644 --- a/controllers/api.go +++ b/controllers/api.go @@ -10,6 +10,7 @@ import ( "text/template" "time" + "github.com/PuerkitoBio/goquery" ctx "github.com/gorilla/context" "github.com/gorilla/mux" "github.com/jinzhu/gorm" @@ -393,12 +394,22 @@ func API_Import_Site(w http.ResponseWriter, r *http.Request) { JSONResponse(w, models.Response{Success: false, Message: err.Error()}, http.StatusBadRequest) return } - body, err := ioutil.ReadAll(resp.Body) + // Insert the base href tag to better handle relative resources + d, err := goquery.NewDocumentFromReader(resp.Body) if err != nil { JSONResponse(w, models.Response{Success: false, Message: err.Error()}, http.StatusBadRequest) return } - cs := cloneResponse{HTML: string(body)} + // Assuming we don't want to include resources, we'll need a base href + if d.Find("head base").Length() == 0 { + d.Find("head").AppendHtml(fmt.Sprintf("", cr.URL)) + } + h, err := d.Html() + if err != nil { + JSONResponse(w, models.Response{Success: false, Message: err.Error()}, http.StatusInternalServerError) + return + } + cs := cloneResponse{HTML: h} JSONResponse(w, cs, http.StatusOK) return } diff --git a/controllers/api_test.go b/controllers/api_test.go new file mode 100644 index 00000000..79169704 --- /dev/null +++ b/controllers/api_test.go @@ -0,0 +1,70 @@ +package controllers + +import ( + "bytes" + "fmt" + "io/ioutil" + "net/http" + "net/http/httptest" + "os" + "testing" + + "github.com/gorilla/handlers" + "github.com/jordan-wright/gophish/config" + "github.com/jordan-wright/gophish/models" + "github.com/stretchr/testify/suite" +) + +// ControllersSuite is a suite of tests to cover API related functions +type ControllersSuite struct { + suite.Suite + ApiKey string +} + +// as is the Admin Server for our API calls +var as *httptest.Server = httptest.NewUnstartedServer(handlers.CombinedLoggingHandler(os.Stdout, CreateAdminRouter())) + +func (s *ControllersSuite) SetupSuite() { + config.Conf.DBPath = ":memory:" + err := models.Setup() + if err != nil { + s.T().Fatalf("Failed creating database: %v", err) + } + s.Nil(err) + // Setup the admin server for use in testing + as.Config.Addr = config.Conf.AdminURL + as.Start() + // Get the API key to use for these tests + u, err := models.GetUser(1) + s.Nil(err) + s.ApiKey = u.ApiKey +} + +func (s *ControllersSuite) TestSiteImportBaseHref() { + h := "" + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprintln(w, h) + })) + defer ts.Close() + resp, err := http.Post(fmt.Sprintf("%s/api/import/site?api_key=%s", as.URL, s.ApiKey), "application/json", + bytes.NewBuffer([]byte(fmt.Sprintf(` + { + "url" : "%s", + "include_resources" : false + } + `, ts.URL)))) + s.Nil(err) + defer resp.Body.Close() + body, err := ioutil.ReadAll(resp.Body) + s.Nil(err) + fmt.Printf("%s", body) +} + +func (s *ControllersSuite) TearDownSuite() { + // Tear down the admin server + as.Close() +} + +func TestControllerSuite(t *testing.T) { + suite.Run(t, new(ControllersSuite)) +}