Added /api/import/site functionality

pull/64/head
Jordan Wright 2015-08-23 20:42:47 -05:00
parent 1ec08d86cf
commit 54fe866d71
2 changed files with 83 additions and 2 deletions

View File

@ -10,6 +10,7 @@ import (
"text/template" "text/template"
"time" "time"
"github.com/PuerkitoBio/goquery"
ctx "github.com/gorilla/context" ctx "github.com/gorilla/context"
"github.com/gorilla/mux" "github.com/gorilla/mux"
"github.com/jinzhu/gorm" "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) JSONResponse(w, models.Response{Success: false, Message: err.Error()}, http.StatusBadRequest)
return 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 { if err != nil {
JSONResponse(w, models.Response{Success: false, Message: err.Error()}, http.StatusBadRequest) JSONResponse(w, models.Response{Success: false, Message: err.Error()}, http.StatusBadRequest)
return 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("<base href=\"%s\">", 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) JSONResponse(w, cs, http.StatusOK)
return return
} }

70
controllers/api_test.go Normal file
View File

@ -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 := "<html><head></head><body><img src=\"/test.png\"/></body></html>"
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))
}