bugfixes, etc, this is super alpha branch so your patch notes are the diff

README.md

@@ -38,13 +38,7 @@ Then you can run puppy by running `puppy`. It will use the puppy binary in `~/$G
 
 Missing Features From Pappy
 ---------------------------
-Here's what Pappy can do that this can't:
-
-- The `http://pappy` interface
-- Upstream proxies
-- Commands taking multiple requests
-- Any and all documentation
-- The macro API is totally different
+All that's left is updating documentation!
 
 Need more info?
 ---------------

certs.go

@@ -3,9 +3,9 @@ package main
 import (
 	"crypto/rand"
 	"crypto/rsa"
+	"crypto/sha1"
 	"crypto/x509"
 	"crypto/x509/pkix"
-	"crypto/sha1"
 	"encoding/pem"
 	"fmt"
 	"math/big"
@@ -14,7 +14,7 @@ import (
 
 type CAKeyPair struct {
 	Certificate []byte
-	PrivateKey *rsa.PrivateKey
+	PrivateKey  *rsa.PrivateKey
 }
 
 func bigIntHash(n *big.Int) []byte {
@@ -41,16 +41,16 @@ func GenerateCACerts() (*CAKeyPair, error) {
 	template := x509.Certificate{
 		SerialNumber: serial,
 		Subject: pkix.Name{
-			CommonName: "Puppy Proxy",
+			CommonName:   "Puppy Proxy",
 			Organization: []string{"Puppy Proxy"},
 		},
 		NotBefore: time.Now().Add(-5 * time.Minute).UTC(),
 		NotAfter:  end,
 
-		SubjectKeyId: bigIntHash(key.N),
-		KeyUsage: x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
+		SubjectKeyId:          bigIntHash(key.N),
+		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
 		BasicConstraintsValid: true,
-		IsCA: true,
+		IsCA:           true,
 		MaxPathLenZero: true,
 	}
 
@@ -61,23 +61,23 @@ func GenerateCACerts() (*CAKeyPair, error) {
 
 	return &CAKeyPair{
 		Certificate: derBytes,
-		PrivateKey: key,
+		PrivateKey:  key,
 	}, nil
 }
 
-func (pair *CAKeyPair) PrivateKeyPEM() ([]byte) {
+func (pair *CAKeyPair) PrivateKeyPEM() []byte {
 	return pem.EncodeToMemory(
 		&pem.Block{
-			Type: "BEGIN PRIVATE KEY",
+			Type:  "BEGIN PRIVATE KEY",
 			Bytes: x509.MarshalPKCS1PrivateKey(pair.PrivateKey),
 		},
 	)
 }
 
-func (pair *CAKeyPair) CACertPEM() ([]byte) {
+func (pair *CAKeyPair) CACertPEM() []byte {
 	return pem.EncodeToMemory(
 		&pem.Block{
-			Type: "CERTIFICATE",
+			Type:  "CERTIFICATE",
 			Bytes: pair.Certificate,
 		},
 	)

credits.go

@@ -5,22 +5,22 @@ List of info that is used to display credits
 */
 
 type creditItem struct {
-	projectName    string
-	url            string
-	author         string
-	year           string
-	licenseType    string
-	longCopyright  string
+	projectName   string
+	url           string
+	author        string
+	year          string
+	licenseType   string
+	longCopyright string
 }
 
-var LIB_CREDITS = []creditItem {
-	creditItem {
+var LIB_CREDITS = []creditItem{
+	creditItem{
 		"goproxy",
 		"https://github.com/elazarl/goproxy",
 		"Elazar Leibovich",
 		"2012",
 		"3-Clause BSD",
-`Copyright (c) 2012 Elazar Leibovich. All rights reserved.
+		`Copyright (c) 2012 Elazar Leibovich. All rights reserved.
 
 Redistribution and use in source and binary forms, with or without
 modification, are permitted provided that the following conditions are
@@ -49,13 +49,13 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.`,
 	},
 
-	creditItem {
+	creditItem{
 		"golang-set",
 		"https://github.com/deckarep/golang-set",
 		"Ralph Caraveo",
 		"2013",
 		"MIT",
-`Open Source Initiative OSI - The MIT License (MIT):Licensing
+		`Open Source Initiative OSI - The MIT License (MIT):Licensing
 
 The MIT License (MIT)
 Copyright (c) 2013 Ralph Caraveo (deckarep@gmail.com)
@@ -79,13 +79,13 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.`,
 	},
 
-	creditItem {
+	creditItem{
 		"Gorilla WebSocket",
 		"https://github.com/gorilla/websocket",
 		"Gorilla WebSocket Authors",
 		"2013",
 		"2-Clause BSD",
-`Copyright (c) 2013 The Gorilla WebSocket Authors. All rights reserved.
+		`Copyright (c) 2013 The Gorilla WebSocket Authors. All rights reserved.
 
 Redistribution and use in source and binary forms, with or without
 modification, are permitted provided that the following conditions are met:

jobpool.go

@@ -7,8 +7,8 @@ import (
 
 // An interface which represents a job to be done by the pool
 type Job interface {
-	Run() // Start the job
-	Abort() // Abort any work that needs to be completed and close the DoneChannel
+	Run()                       // Start the job
+	Abort()                     // Abort any work that needs to be completed and close the DoneChannel
 	DoneChannel() chan struct{} // Returns a channel that is closed when the job is done
 }
 
@@ -16,19 +16,19 @@ type Job interface {
 type JobPool struct {
 	MaxThreads int
 
-	jobQueue chan Job
-	jobQueueDone chan struct{}
-	jobQueueAborted chan struct{}
+	jobQueue         chan Job
+	jobQueueDone     chan struct{}
+	jobQueueAborted  chan struct{}
 	jobQueueShutDown chan struct{}
 }
 
-func NewJobPool(maxThreads int) (*JobPool) {
-	q := JobPool {
-		MaxThreads: maxThreads,
-		jobQueue: make(chan Job),
-		jobQueueDone: make(chan struct{}),      // Closing will shut down workers and reject any incoming work
-		jobQueueAborted: make(chan struct{}),   // Closing tells workers to abort
-		jobQueueShutDown: make(chan struct{}),  // Closed when all workers are shut down
+func NewJobPool(maxThreads int) *JobPool {
+	q := JobPool{
+		MaxThreads:       maxThreads,
+		jobQueue:         make(chan Job),
+		jobQueueDone:     make(chan struct{}), // Closing will shut down workers and reject any incoming work
+		jobQueueAborted:  make(chan struct{}), // Closing tells workers to abort
+		jobQueueShutDown: make(chan struct{}), // Closed when all workers are shut down
 	}
 
 	return &q
@@ -49,7 +49,7 @@ func (q *JobPool) Run() {
 	if q.MaxThreads > 0 {
 		// Create pool of routines that read from the queue and run jobs
 		var w sync.WaitGroup
-		for i:=0; i<q.MaxThreads; i++ {
+		for i := 0; i < q.MaxThreads; i++ {
 			w.Add(1)
 			go func() {
 				defer w.Done()
@@ -75,7 +75,7 @@ func (q *JobPool) Run() {
 			}()
 		}
 
-		w.Wait() // Wait for workers to quit
+		w.Wait()                  // Wait for workers to quit
 		close(q.jobQueueShutDown) // Flag that all workers quit
 	} else {
 		// Create a thread any time we pull something out of the job queue
@@ -103,15 +103,15 @@ func (q *JobPool) Run() {
 }
 
 func (q *JobPool) Abort() {
-	close(q.jobQueueDone) // Stop accepting jobs and tell the workers to quit
+	close(q.jobQueueDone)    // Stop accepting jobs and tell the workers to quit
 	close(q.jobQueueAborted) // Tell the workers to abort
-	<-q.jobQueueShutDown  // Wait for all the workers to shut down
-	close(q.jobQueue)     // Clean up the job queue
+	<-q.jobQueueShutDown     // Wait for all the workers to shut down
+	close(q.jobQueue)        // Clean up the job queue
 }
 
 func (q *JobPool) CompleteAndClose() {
-	close(q.jobQueueDone) // Stop accepting jobs and tell the workers to quit
-	<-q.jobQueueShutDown  // Wait for all the workers to shut down
-	close(q.jobQueue)     // Clean up the job queue
+	close(q.jobQueueDone)    // Stop accepting jobs and tell the workers to quit
+	<-q.jobQueueShutDown     // Wait for all the workers to shut down
+	close(q.jobQueue)        // Clean up the job queue
 	close(q.jobQueueAborted) // Clean up abort channel
 }

main.go

@@ -6,11 +6,11 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
-	"strings"
-	"syscall"
-	"os/signal"
 	"net"
 	"os"
+	"os/signal"
+	"strings"
+	"syscall"
 	"time"
 )
 
@@ -101,7 +101,7 @@ func main() {
 	flag.Parse()
 
 	if *debugFlag {
-		logfile, err := os.OpenFile("log.log", os.O_RDWR | os.O_CREATE | os.O_APPEND, 0666)
+		logfile, err := os.OpenFile("log.log", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
 		checkErr(err)
 		logger = log.New(logfile, "[*] ", log.Lshortfile)
 	} else {
@@ -151,28 +151,9 @@ func main() {
 		}
 	}
 
-	// Set up storage
-	// if *storageFname == "" && *inMemStorage == false {
-	// 	quitErr("storage file or -inmem flag required")
-	// }
-	// if *storageFname != "" && *inMemStorage == true {
-	// 	quitErr("cannot provide both a storage file and -inmem flag")
-	// }
-	// var storage MessageStorage
-	// if *inMemStorage {
-	// 	var err error
-	// 	storage, err = InMemoryStorage(logger)
-	// 	checkErr(err)
-	// } else {
-	// 	var err error
-	// 	storage, err = OpenSQLiteStorage(*storageFname, logger)
-	// 	checkErr(err)
-	// }
-
 	// Set up the intercepting proxy
 	iproxy := NewInterceptingProxy(logger)
-	// sid := iproxy.AddMessageStorage(storage)
-	// iproxy.SetProxyStorage(sid)
+	iproxy.AddHTTPHandler("puppy", WebUIHandler)
 
 	// Create a message server and have it serve for the iproxy
 	mserv := NewProxyMessageListener(logger, iproxy)

messageserv.go

@@ -4,10 +4,10 @@ import (
 	"bufio"
 	"encoding/json"
 	"fmt"
-	"strings"
 	"io"
 	"log"
 	"net"
+	"strings"
 )
 
 /*
@@ -18,8 +18,8 @@ type MessageHandler func([]byte, net.Conn, *log.Logger, *InterceptingProxy)
 
 type MessageListener struct {
 	handlers map[string]MessageHandler
-	iproxy *InterceptingProxy
-	Logger *log.Logger
+	iproxy   *InterceptingProxy
+	Logger   *log.Logger
 }
 
 type commandData struct {
@@ -28,14 +28,14 @@ type commandData struct {
 
 type errorMessage struct {
 	Success bool
-	Reason string
+	Reason  string
 }
 
 func NewMessageListener(l *log.Logger, iproxy *InterceptingProxy) *MessageListener {
 	m := &MessageListener{
 		handlers: make(map[string]MessageHandler),
-		iproxy: iproxy,
-		Logger: l,
+		iproxy:   iproxy,
+		Logger:   l,
 	}
 	return m
 }

proxy.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"crypto/tls"
+	"encoding/base64"
 	"fmt"
 	"log"
 	"net"
@@ -15,61 +16,80 @@ import (
 var getNextSubId = IdCounter()
 var getNextStorageId = IdCounter()
 
+// Working on using this for webui
+type proxyWebUIHandler func(http.ResponseWriter, *http.Request, *InterceptingProxy)
+
 type savedStorage struct {
-	storage MessageStorage
+	storage     MessageStorage
 	description string
 }
 
 type InterceptingProxy struct {
-	slistener *ProxyListener
-	server    *http.Server
-	mtx       sync.Mutex
-	logger    *log.Logger
+	slistener    *ProxyListener
+	server       *http.Server
+	mtx          sync.Mutex
+	logger       *log.Logger
 	proxyStorage int
+	netDial      NetDialer
+
+	usingProxy   bool
+	proxyHost    string
+	proxyPort    int
+	proxyIsSOCKS bool
+	proxyCreds   *ProxyCredentials
 
-	requestInterceptor RequestInterceptor
+	requestInterceptor  RequestInterceptor
 	responseInterceptor ResponseInterceptor
-	wSInterceptor WSInterceptor
-	scopeChecker RequestChecker
-	scopeQuery MessageQuery
+	wSInterceptor       WSInterceptor
+	scopeChecker        RequestChecker
+	scopeQuery          MessageQuery
 
 	reqSubs []*ReqIntSub
 	rspSubs []*RspIntSub
-	wsSubs []*WSIntSub
+	wsSubs  []*WSIntSub
+
+	httpHandlers map[string]proxyWebUIHandler
 
 	messageStorage map[int]*savedStorage
 }
 
+type ProxyCredentials struct {
+	Username string
+	Password string
+}
+
 type RequestInterceptor func(req *ProxyRequest) (*ProxyRequest, error)
 type ResponseInterceptor func(req *ProxyRequest, rsp *ProxyResponse) (*ProxyResponse, error)
 type WSInterceptor func(req *ProxyRequest, rsp *ProxyResponse, msg *ProxyWSMessage) (*ProxyWSMessage, error)
 
-type proxyHandler struct {
-	Logger *log.Logger
-	IProxy *InterceptingProxy
-}
-
 type ReqIntSub struct {
-	id int
+	id          int
 	Interceptor RequestInterceptor
 }
 
 type RspIntSub struct {
-	id int
+	id          int
 	Interceptor ResponseInterceptor
 }
 
 type WSIntSub struct {
-	id int
+	id          int
 	Interceptor WSInterceptor
 }
 
+func (creds *ProxyCredentials) SerializeHeader() string {
+	toEncode := []byte(fmt.Sprintf("%s:%s", creds.Username, creds.Password))
+	encoded := base64.StdEncoding.EncodeToString(toEncode)
+	return fmt.Sprintf("Basic %s", encoded)
+}
+
 func NewInterceptingProxy(logger *log.Logger) *InterceptingProxy {
 	var iproxy InterceptingProxy
 	iproxy.messageStorage = make(map[int]*savedStorage)
 	iproxy.slistener = NewProxyListener(logger)
 	iproxy.server = newProxyServer(logger, &iproxy)
 	iproxy.logger = logger
+	iproxy.httpHandlers = make(map[string]proxyWebUIHandler)
 
 	go func() {
 		iproxy.server.Serve(iproxy.slistener)
@@ -93,7 +113,7 @@ func (iproxy *InterceptingProxy) SetCACertificate(caCert *tls.Certificate) {
 	iproxy.slistener.SetCACertificate(caCert)
 }
 
-func (iproxy *InterceptingProxy) GetCACertificate() (*tls.Certificate) {
+func (iproxy *InterceptingProxy) GetCACertificate() *tls.Certificate {
 	return iproxy.slistener.GetCACertificate()
 }
 
@@ -119,7 +139,7 @@ func (iproxy *InterceptingProxy) GetMessageStorage(id int) (MessageStorage, stri
 	return savedStorage.storage, savedStorage.description
 }
 
-func (iproxy *InterceptingProxy) AddMessageStorage(storage MessageStorage, description string) (int) {
+func (iproxy *InterceptingProxy) AddMessageStorage(storage MessageStorage, description string) int {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 	id := getNextStorageId()
@@ -139,8 +159,8 @@ func (iproxy *InterceptingProxy) CloseMessageStorage(id int) {
 }
 
 type SavedStorage struct {
-	Id int
-	Storage MessageStorage
+	Id          int
+	Storage     MessageStorage
 	Description string
 }
 
@@ -190,7 +210,7 @@ func (iproxy *InterceptingProxy) LoadScope(storageId int) error {
 	return nil
 }
 
-func (iproxy *InterceptingProxy) GetScopeChecker() (RequestChecker) {
+func (iproxy *InterceptingProxy) GetScopeChecker() RequestChecker {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 	return iproxy.scopeChecker
@@ -212,19 +232,19 @@ func (iproxy *InterceptingProxy) SetScopeChecker(checker RequestChecker) error {
 	return nil
 }
 
-func (iproxy *InterceptingProxy) GetScopeQuery() (MessageQuery) {
+func (iproxy *InterceptingProxy) GetScopeQuery() MessageQuery {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 	return iproxy.scopeQuery
 }
 
-func (iproxy *InterceptingProxy) SetScopeQuery(query MessageQuery) (error) {
+func (iproxy *InterceptingProxy) SetScopeQuery(query MessageQuery) error {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 	return iproxy.setScopeQuery(query)
 }
 
-func (iproxy *InterceptingProxy) setScopeQuery(query MessageQuery) (error) {
+func (iproxy *InterceptingProxy) setScopeQuery(query MessageQuery) error {
 	checker, err := CheckerFromMessageQuery(query)
 	if err != nil {
 		return err
@@ -244,7 +264,48 @@ func (iproxy *InterceptingProxy) setScopeQuery(query MessageQuery) (error) {
 	return nil
 }
 
-func (iproxy *InterceptingProxy) ClearScope() (error) {
+func (iproxy *InterceptingProxy) SetNetDial(dialer NetDialer) {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	iproxy.netDial = dialer
+}
+
+func (iproxy *InterceptingProxy) NetDial() NetDialer {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	return iproxy.netDial
+}
+
+func (iproxy *InterceptingProxy) ClearUpstreamProxy() {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	iproxy.usingProxy = false
+	iproxy.proxyHost = ""
+	iproxy.proxyPort = 0
+	iproxy.proxyIsSOCKS = false
+}
+
+func (iproxy *InterceptingProxy) SetUpstreamProxy(proxyHost string, proxyPort int, creds *ProxyCredentials) {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	iproxy.usingProxy = true
+	iproxy.proxyHost = proxyHost
+	iproxy.proxyPort = proxyPort
+	iproxy.proxyIsSOCKS = false
+	iproxy.proxyCreds = creds
+}
+
+func (iproxy *InterceptingProxy) SetUpstreamSOCKSProxy(proxyHost string, proxyPort int, creds *ProxyCredentials) {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	iproxy.usingProxy = true
+	iproxy.proxyHost = proxyHost
+	iproxy.proxyPort = proxyPort
+	iproxy.proxyIsSOCKS = true
+	iproxy.proxyCreds = creds
+}
+
+func (iproxy *InterceptingProxy) ClearScope() error {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 	iproxy.scopeChecker = nil
@@ -262,12 +323,42 @@ func (iproxy *InterceptingProxy) ClearScope() (error) {
 	return nil
 }
 
-func (iproxy *InterceptingProxy) AddReqIntSub(f RequestInterceptor) (*ReqIntSub) {
+func (iproxy *InterceptingProxy) SubmitRequest(req *ProxyRequest) error {
+	oldDial := req.NetDial
+	defer func() { req.NetDial = oldDial }()
+	req.NetDial = iproxy.NetDial()
+
+	if iproxy.usingProxy {
+		if iproxy.proxyIsSOCKS {
+			return SubmitRequestSOCKSProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds)
+		} else {
+			return SubmitRequestProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds)
+		}
+	}
+	return SubmitRequest(req)
+}
+
+func (iproxy *InterceptingProxy) WSDial(req *ProxyRequest) (*WSSession, error) {
+	oldDial := req.NetDial
+	defer func() { req.NetDial = oldDial }()
+	req.NetDial = iproxy.NetDial()
+
+	if iproxy.usingProxy {
+		if iproxy.proxyIsSOCKS {
+			return WSDialSOCKSProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds)
+		} else {
+			return WSDialProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds)
+		}
+	}
+	return WSDial(req)
+}
+
+func (iproxy *InterceptingProxy) AddReqIntSub(f RequestInterceptor) *ReqIntSub {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 
 	sub := &ReqIntSub{
-		id: getNextSubId(),
+		id:          getNextSubId(),
 		Interceptor: f,
 	}
 	iproxy.reqSubs = append(iproxy.reqSubs, sub)
@@ -286,12 +377,12 @@ func (iproxy *InterceptingProxy) RemoveReqIntSub(sub *ReqIntSub) {
 	}
 }
 
-func (iproxy *InterceptingProxy) AddRspIntSub(f ResponseInterceptor) (*RspIntSub) {
+func (iproxy *InterceptingProxy) AddRspIntSub(f ResponseInterceptor) *RspIntSub {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 
 	sub := &RspIntSub{
-		id: getNextSubId(),
+		id:          getNextSubId(),
 		Interceptor: f,
 	}
 	iproxy.rspSubs = append(iproxy.rspSubs, sub)
@@ -310,12 +401,12 @@ func (iproxy *InterceptingProxy) RemoveRspIntSub(sub *RspIntSub) {
 	}
 }
 
-func (iproxy *InterceptingProxy) AddWSIntSub(f WSInterceptor) (*WSIntSub) {
+func (iproxy *InterceptingProxy) AddWSIntSub(f WSInterceptor) *WSIntSub {
 	iproxy.mtx.Lock()
 	defer iproxy.mtx.Unlock()
 
 	sub := &WSIntSub{
-		id: getNextSubId(),
+		id:          getNextSubId(),
 		Interceptor: f,
 	}
 	iproxy.wsSubs = append(iproxy.wsSubs, sub)
@@ -360,6 +451,28 @@ func (iproxy *InterceptingProxy) GetProxyStorage() MessageStorage {
 	return savedStorage.storage
 }
 
+func (iproxy *InterceptingProxy) AddHTTPHandler(host string, handler proxyWebUIHandler) {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	iproxy.httpHandlers[host] = handler
+}
+
+func (iproxy *InterceptingProxy) GetHTTPHandler(host string) (proxyWebUIHandler, error) {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	handler, ok := iproxy.httpHandlers[host]
+	if !ok {
+		return nil, fmt.Errorf("no handler for host %s", host)
+	}
+	return handler, nil
+}
+
+func (iproxy *InterceptingProxy) RemoveHTTPHandler(host string) {
+	iproxy.mtx.Lock()
+	defer iproxy.mtx.Unlock()
+	delete(iproxy.httpHandlers, host)
+}
+
 func ParseProxyRequest(r *http.Request) (*ProxyRequest, error) {
 	host, port, useTLS, err := DecodeRemoteAddr(r.RemoteAddr)
 	if err != nil {
@@ -382,14 +495,19 @@ func ErrResponse(w http.ResponseWriter, err error) {
 	http.Error(w, err.Error(), http.StatusInternalServerError)
 }
 
-func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+func (iproxy *InterceptingProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	handler, err := iproxy.GetHTTPHandler(r.Host)
+	if err == nil {
+		handler(w, r, iproxy)
+		return
+	}
 
 	req, _ := ParseProxyRequest(r)
-	p.Logger.Println("Received request to", req.FullURL().String())
+	iproxy.logger.Println("Received request to", req.FullURL().String())
 	req.StripProxyHeaders()
 
-	ms := p.IProxy.GetProxyStorage()
-	scopeChecker := p.IProxy.GetScopeChecker()
+	ms := iproxy.GetProxyStorage()
+	scopeChecker := iproxy.GetScopeChecker()
 
 	// Helper functions
 	checkScope := func(req *ProxyRequest) bool {
@@ -409,16 +527,16 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	/*
-	functions to mangle messages using the iproxy's manglers
-	each return the new message, whether it was modified, and an error
+		functions to mangle messages using the iproxy's manglers
+		each return the new message, whether it was modified, and an error
 	*/
 
 	mangleRequest := func(req *ProxyRequest) (*ProxyRequest, bool, error) {
 		newReq := req.Clone()
-		reqSubs := p.IProxy.getRequestSubs()
+		reqSubs := iproxy.getRequestSubs()
 		for _, sub := range reqSubs {
 			var err error = nil
-			newReq, err := sub.Interceptor(newReq)
+			newReq, err = sub.Interceptor(newReq)
 			if err != nil {
 				e := fmt.Errorf("error with request interceptor: %s", err)
 				return nil, false, e
@@ -431,7 +549,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		if newReq != nil {
 			newReq.StartDatetime = time.Now()
 			if !req.Eq(newReq) {
-				p.Logger.Println("Request modified by interceptor")
+				iproxy.logger.Println("Request modified by interceptor")
 				return newReq, true, nil
 			}
 		} else {
@@ -443,10 +561,10 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	mangleResponse := func(req *ProxyRequest, rsp *ProxyResponse) (*ProxyResponse, bool, error) {
 		reqCopy := req.Clone()
 		newRsp := rsp.Clone()
-		rspSubs := p.IProxy.getResponseSubs()
-		p.Logger.Printf("%d interceptors", len(rspSubs))
+		rspSubs := iproxy.getResponseSubs()
+		iproxy.logger.Printf("%d interceptors", len(rspSubs))
 		for _, sub := range rspSubs {
-			p.Logger.Println("mangling rsp...")
+			iproxy.logger.Println("mangling rsp...")
 			var err error = nil
 			newRsp, err = sub.Interceptor(reqCopy, newRsp)
 			if err != nil {
@@ -460,7 +578,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 		if newRsp != nil {
 			if !rsp.Eq(newRsp) {
-				p.Logger.Println("Response for", req.FullURL(), "modified by interceptor")
+				iproxy.logger.Println("Response for", req.FullURL(), "modified by interceptor")
 				// it was mangled
 				return newRsp, true, nil
 			}
@@ -477,7 +595,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		newMsg := ws.Clone()
 		reqCopy := req.Clone()
 		rspCopy := rsp.Clone()
-		wsSubs := p.IProxy.getWSSubs()
+		wsSubs := iproxy.getWSSubs()
 		for _, sub := range wsSubs {
 			var err error = nil
 			newMsg, err = sub.Interceptor(reqCopy, rspCopy, newMsg)
@@ -494,7 +612,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			if !ws.Eq(newMsg) {
 				newMsg.Timestamp = time.Now()
 				newMsg.Direction = ws.Direction
-				p.Logger.Println("Message modified by interceptor")
+				iproxy.logger.Println("Message modified by interceptor")
 				return newMsg, true, nil
 			}
 		} else {
@@ -503,7 +621,6 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return ws, false, nil
 	}
 
-
 	req.StartDatetime = time.Now()
 
 	if checkScope(req) {
@@ -537,12 +654,12 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if req.IsWSUpgrade() {
-		p.Logger.Println("Detected websocket request. Upgrading...")
+		iproxy.logger.Println("Detected websocket request. Upgrading...")
 
-		rc, err := req.WSDial()
+		rc, err := iproxy.WSDial(req)
 		if err != nil {
-			p.Logger.Println("error dialing ws server:", err)
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+			iproxy.logger.Println("error dialing ws server:", err)
+			http.Error(w, fmt.Sprintf("error dialing websocket server: %s", err.Error()), http.StatusInternalServerError)
 			return
 		}
 		defer rc.Close()
@@ -560,8 +677,8 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 		lc, err := upgrader.Upgrade(w, r, nil)
 		if err != nil {
-			p.Logger.Println("error upgrading connection:", err)
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+			iproxy.logger.Println("error upgrading connection:", err)
+			http.Error(w, fmt.Sprintf("error upgrading connection: %s", err.Error()), http.StatusInternalServerError)
 			return
 		}
 		defer lc.Close()
@@ -581,13 +698,13 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 				mtype, msg, err := rc.ReadMessage()
 				if err != nil {
 					lc.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
-					p.Logger.Println("error with receiving server message:", err)
+					iproxy.logger.Println("error with receiving server message:", err)
 					wg.Done()
 					return
 				}
 				pws, err := NewProxyWSMessage(mtype, msg, ToClient)
 				if err != nil {
-					p.Logger.Println("error creating ws object:", err.Error())
+					iproxy.logger.Println("error creating ws object:", err.Error())
 					continue
 				}
 				pws.Timestamp = time.Now()
@@ -595,7 +712,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 				if checkScope(req) {
 					newMsg, mangled, err := mangleWS(req, req.ServerResponse, pws)
 					if err != nil {
-						p.Logger.Println("error mangling ws:", err)
+						iproxy.logger.Println("error mangling ws:", err)
 						return
 					}
 					if mangled {
@@ -611,7 +728,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 				addWSMessage(req, pws)
 				if err := saveIfExists(req); err != nil {
-					p.Logger.Println("error saving request:", err)
+					iproxy.logger.Println("error saving request:", err)
 					continue
 				}
 				lc.WriteMessage(pws.Type, pws.Message)
@@ -625,13 +742,13 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 				mtype, msg, err := lc.ReadMessage()
 				if err != nil {
 					rc.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
-					p.Logger.Println("error with receiving client message:", err)
+					iproxy.logger.Println("error with receiving client message:", err)
 					wg.Done()
 					return
 				}
 				pws, err := NewProxyWSMessage(mtype, msg, ToServer)
 				if err != nil {
-					p.Logger.Println("error creating ws object:", err.Error())
+					iproxy.logger.Println("error creating ws object:", err.Error())
 					continue
 				}
 				pws.Timestamp = time.Now()
@@ -639,7 +756,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 				if checkScope(req) {
 					newMsg, mangled, err := mangleWS(req, req.ServerResponse, pws)
 					if err != nil {
-						p.Logger.Println("error mangling ws:", err)
+						iproxy.logger.Println("error mangling ws:", err)
 						return
 					}
 					if mangled {
@@ -655,18 +772,18 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 				addWSMessage(req, pws)
 				if err := saveIfExists(req); err != nil {
-					p.Logger.Println("error saving request:", err)
+					iproxy.logger.Println("error saving request:", err)
 					continue
 				}
 				rc.WriteMessage(pws.Type, pws.Message)
 			}
 		}()
 		wg.Wait()
-		p.Logger.Println("Websocket session complete!")
+		iproxy.logger.Println("Websocket session complete!")
 	} else {
-		err := req.Submit()
+		err := iproxy.SubmitRequest(req)
 		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+			http.Error(w, fmt.Sprintf("error submitting request: %s", err.Error()), http.StatusInternalServerError)
 			return
 		}
 		req.EndDatetime = time.Now()
@@ -713,10 +830,7 @@ func (p proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 func newProxyServer(logger *log.Logger, iproxy *InterceptingProxy) *http.Server {
 	server := &http.Server{
-		Handler: proxyHandler{
-			Logger: logger,
-			IProxy: iproxy,
-		},
+		Handler:  iproxy,
 		ErrorLog: logger,
 	}
 	return server

proxyhttp.go

@@ -15,12 +15,13 @@ import (
 	"net/http"
 	"net/url"
 	"reflect"
-	"strings"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/deckarep/golang-set"
 	"github.com/gorilla/websocket"
+	"golang.org/x/net/proxy"
 )
 
 const (
@@ -28,10 +29,12 @@ const (
 	ToClient
 )
 
+type NetDialer func(network, addr string) (net.Conn, error)
+
 type ProxyResponse struct {
 	http.Response
 	bodyBytes []byte
-	DbId string // ID used by storage implementation. Blank string = unsaved
+	DbId      string // ID used by storage implementation. Blank string = unsaved
 	Unmangled *ProxyResponse
 }
 
@@ -49,12 +52,14 @@ type ProxyRequest struct {
 	Unmangled      *ProxyRequest
 
 	// Additional data
-	bodyBytes []byte
-	DbId      string // ID used by storage implementation. Blank string = unsaved
+	bodyBytes     []byte
+	DbId          string // ID used by storage implementation. Blank string = unsaved
 	StartDatetime time.Time
 	EndDatetime   time.Time
 
 	tags mapset.Set
+
+	NetDial NetDialer
 }
 
 type WSSession struct {
@@ -64,17 +69,30 @@ type WSSession struct {
 }
 
 type ProxyWSMessage struct {
-	Type int
-	Message []byte
+	Type      int
+	Message   []byte
 	Direction int
 	Unmangled *ProxyWSMessage
 	Timestamp time.Time
-	Request *ProxyRequest
+	Request   *ProxyRequest
 
 	DbId string // ID used by storage implementation. Blank string = unsaved
 }
 
-func NewProxyRequest(r *http.Request, destHost string, destPort int, destUseTLS bool) (*ProxyRequest) {
+func PerformConnect(conn net.Conn, destHost string, destPort int) error {
+	connStr := []byte(fmt.Sprintf("CONNECT %s:%d HTTP/1.1\r\nHost: %s\r\nProxy-Connection: Keep-Alive\r\n\r\n", destHost, destPort, destHost))
+	conn.Write(connStr)
+	rsp, err := http.ReadResponse(bufio.NewReader(conn), nil)
+	if err != nil {
+		return fmt.Errorf("error performing CONNECT handshake: %s", err.Error())
+	}
+	if rsp.StatusCode != 200 {
+		return fmt.Errorf("error performing CONNECT handshake")
+	}
+	return nil
+}
+
+func NewProxyRequest(r *http.Request, destHost string, destPort int, destUseTLS bool) *ProxyRequest {
 	var retReq *ProxyRequest
 	if r != nil {
 		// Write/reread the request to make sure we get all the extra headers Go adds into req.Header
@@ -98,6 +116,7 @@ func NewProxyRequest(r *http.Request, destHost string, destPort int, destUseTLS
 			time.Unix(0, 0),
 			time.Unix(0, 0),
 			mapset.NewSet(),
+			nil,
 		}
 	} else {
 		newReq, _ := http.NewRequest("GET", "/", nil) // Ignore error since this should be run the same every time and shouldn't error
@@ -116,6 +135,7 @@ func NewProxyRequest(r *http.Request, destHost string, destPort int, destUseTLS
 			time.Unix(0, 0),
 			time.Unix(0, 0),
 			mapset.NewSet(),
+			nil,
 		}
 	}
 
@@ -125,7 +145,7 @@ func NewProxyRequest(r *http.Request, destHost string, destPort int, destUseTLS
 	return retReq
 }
 
-func ProxyRequestFromBytes(b []byte,  destHost string, destPort int, destUseTLS bool) (*ProxyRequest, error) {
+func ProxyRequestFromBytes(b []byte, destHost string, destPort int, destUseTLS bool) (*ProxyRequest, error) {
 	buf := bytes.NewBuffer(b)
 	httpReq, err := http.ReadRequest(bufio.NewReader(buf))
 	if err != nil {
@@ -135,7 +155,7 @@ func ProxyRequestFromBytes(b []byte,  destHost string, destPort int, destUseTLS
 	return NewProxyRequest(httpReq, destHost, destPort, destUseTLS), nil
 }
 
-func NewProxyResponse(r *http.Response) (*ProxyResponse) {
+func NewProxyResponse(r *http.Response) *ProxyResponse {
 	// Write/reread the request to make sure we get all the extra headers Go adds into req.Header
 	oldClose := r.Close
 	r.Close = false
@@ -170,12 +190,12 @@ func ProxyResponseFromBytes(b []byte) (*ProxyResponse, error) {
 
 func NewProxyWSMessage(mtype int, message []byte, direction int) (*ProxyWSMessage, error) {
 	return &ProxyWSMessage{
-		Type: mtype,
-		Message: message,
+		Type:      mtype,
+		Message:   message,
 		Direction: direction,
 		Unmangled: nil,
 		Timestamp: time.Unix(0, 0),
-		DbId: "",
+		DbId:      "",
 	}, nil
 }
 
@@ -221,33 +241,38 @@ func (req *ProxyRequest) DestURL() *url.URL {
 	return &u
 }
 
-func (req *ProxyRequest) Submit() error {
-	// Connect to the remote server
-	var conn net.Conn
-	var err error
-	dest := fmt.Sprintf("%s:%d", req.DestHost, req.DestPort)
-	if req.DestUseTLS {
-		// Use TLS
-		conn, err = tls.Dial("tcp", dest, nil)
-		if err != nil {
+func (req *ProxyRequest) Submit(conn net.Conn) error {
+	return req.submit(conn, false, nil)
+}
+
+func (req *ProxyRequest) SubmitProxy(conn net.Conn, creds *ProxyCredentials) error {
+	return req.submit(conn, true, creds)
+}
+
+func (req *ProxyRequest) submit(conn net.Conn, forProxy bool, proxyCreds *ProxyCredentials) error {
+	// Write the request to the connection
+	req.StartDatetime = time.Now()
+	if forProxy {
+		if req.DestUseTLS {
+			req.URL.Scheme = "https"
+		} else {
+			req.URL.Scheme = "http"
+		}
+		req.URL.Opaque = ""
+
+		if err := req.RepeatableProxyWrite(conn, proxyCreds); err != nil {
 			return err
 		}
 	} else {
-		// Use plaintext
-		conn, err = net.Dial("tcp", dest)
-		if err != nil {
+		if err := req.RepeatableWrite(conn); err != nil {
 			return err
 		}
 	}
 
-	// Write the request to the connection
-	req.StartDatetime = time.Now()
-	req.RepeatableWrite(conn)
-
 	// Read a response from the server
 	httpRsp, err := http.ReadResponse(bufio.NewReader(conn), nil)
 	if err != nil {
-		return err
+		return fmt.Errorf("error reading response: %s", err.Error())
 	}
 	req.EndDatetime = time.Now()
 
@@ -256,7 +281,7 @@ func (req *ProxyRequest) Submit() error {
 	return nil
 }
 
-func (req *ProxyRequest) WSDial() (*WSSession, error) {
+func (req *ProxyRequest) WSDial(conn net.Conn) (*WSSession, error) {
 	if !req.IsWSUpgrade() {
 		return nil, fmt.Errorf("could not start websocket session: request is not a websocket handshake request")
 	}
@@ -276,18 +301,91 @@ func (req *ProxyRequest) WSDial() (*WSSession, error) {
 	}
 
 	dialer := &websocket.Dialer{}
-	conn, rsp, err := dialer.Dial(req.DestURL().String(), upgradeHeaders)
+	dialer.NetDial = func(network, address string) (net.Conn, error) {
+		return conn, nil
+	}
+
+	wsconn, rsp, err := dialer.Dial(req.DestURL().String(), upgradeHeaders)
 	if err != nil {
 		return nil, fmt.Errorf("could not dial WebSocket server: %s", err)
 	}
 	req.ServerResponse = NewProxyResponse(rsp)
 	wsession := &WSSession{
-		*conn,
+		*wsconn,
 		req,
 	}
 	return wsession, nil
 }
 
+func WSDial(req *ProxyRequest) (*WSSession, error) {
+	return wsDial(req, false, "", 0, nil, false)
+}
+
+func WSDialProxy(req *ProxyRequest, proxyHost string, proxyPort int, creds *ProxyCredentials) (*WSSession, error) {
+	return wsDial(req, true, proxyHost, proxyPort, creds, false)
+}
+
+func WSDialSOCKSProxy(req *ProxyRequest, proxyHost string, proxyPort int, creds *ProxyCredentials) (*WSSession, error) {
+	return wsDial(req, true, proxyHost, proxyPort, creds, true)
+}
+
+func wsDial(req *ProxyRequest, useProxy bool, proxyHost string, proxyPort int, proxyCreds *ProxyCredentials, proxyIsSOCKS bool) (*WSSession, error) {
+	var conn net.Conn
+	var dialer NetDialer
+	var err error
+
+	if req.NetDial != nil {
+		dialer = req.NetDial
+	} else {
+		dialer = net.Dial
+	}
+
+	if useProxy {
+		if proxyIsSOCKS {
+			var socksCreds *proxy.Auth
+			if proxyCreds != nil {
+				socksCreds = &proxy.Auth{
+					User:     proxyCreds.Username,
+					Password: proxyCreds.Password,
+				}
+			}
+			socksDialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("%s:%d", proxyHost, proxyPort), socksCreds, proxy.Direct)
+			if err != nil {
+				return nil, fmt.Errorf("error creating SOCKS dialer: %s", err.Error())
+			}
+			conn, err = socksDialer.Dial("tcp", fmt.Sprintf("%s:%d", req.DestHost, req.DestPort))
+			if err != nil {
+				return nil, fmt.Errorf("error dialing host: %s", err.Error())
+			}
+			defer conn.Close()
+		} else {
+			conn, err = dialer("tcp", fmt.Sprintf("%s:%d", proxyHost, proxyPort))
+			if err != nil {
+				return nil, fmt.Errorf("error dialing proxy: %s", err.Error())
+			}
+
+			// always perform a CONNECT for websocket regardless of SSL
+			if err := PerformConnect(conn, req.DestHost, req.DestPort); err != nil {
+				return nil, err
+			}
+		}
+	} else {
+		conn, err = dialer("tcp", fmt.Sprintf("%s:%d", req.DestHost, req.DestPort))
+		if err != nil {
+			return nil, fmt.Errorf("error dialing host: %s", err.Error())
+		}
+	}
+
+	if req.DestUseTLS {
+		tls_conn := tls.Client(conn, &tls.Config{
+			InsecureSkipVerify: true,
+		})
+		conn = tls_conn
+	}
+
+	return req.WSDial(conn)
+}
+
 func (req *ProxyRequest) IsWSUpgrade() bool {
 	for k, v := range req.Header {
 		for _, vv := range v {
@@ -322,7 +420,7 @@ func (req *ProxyRequest) Eq(other *ProxyRequest) bool {
 	return true
 }
 
-func (req *ProxyRequest) Clone() (*ProxyRequest) {
+func (req *ProxyRequest) Clone() *ProxyRequest {
 	buf := bytes.NewBuffer(make([]byte, 0))
 	req.RepeatableWrite(buf)
 	newReq, err := ProxyRequestFromBytes(buf.Bytes(), req.DestHost, req.DestPort, req.DestUseTLS)
@@ -336,7 +434,7 @@ func (req *ProxyRequest) Clone() (*ProxyRequest) {
 	return newReq
 }
 
-func (req *ProxyRequest) DeepClone() (*ProxyRequest) {
+func (req *ProxyRequest) DeepClone() *ProxyRequest {
 	// Returns a request with the same request, response, and associated websocket messages
 	newReq := req.Clone()
 	newReq.DbId = req.DbId
@@ -361,9 +459,19 @@ func (req *ProxyRequest) resetBodyReader() {
 	req.Body = ioutil.NopCloser(bytes.NewBuffer(req.BodyBytes()))
 }
 
-func (req *ProxyRequest) RepeatableWrite(w io.Writer) {
-	req.Write(w)
-	req.resetBodyReader()
+func (req *ProxyRequest) RepeatableWrite(w io.Writer) error {
+	defer req.resetBodyReader()
+	return req.Write(w)
+}
+
+func (req *ProxyRequest) RepeatableProxyWrite(w io.Writer, proxyCreds *ProxyCredentials) error {
+	defer req.resetBodyReader()
+	if proxyCreds != nil {
+		authHeader := proxyCreds.SerializeHeader()
+		req.Header.Set("Proxy-Authorization", authHeader)
+		defer func() { req.Header.Del("Proxy-Authorization") }()
+	}
+	return req.WriteProxy(w)
 }
 
 func (req *ProxyRequest) BodyBytes() []byte {
@@ -376,7 +484,7 @@ func (req *ProxyRequest) SetBodyBytes(bs []byte) {
 	req.resetBodyReader()
 
 	// Parse the form if we can, ignore errors
-	req.ParseMultipartForm(1024*1024*1024) // 1GB for no good reason
+	req.ParseMultipartForm(1024 * 1024 * 1024) // 1GB for no good reason
 	req.ParseForm()
 	req.resetBodyReader()
 	req.Header.Set("Content-Length", strconv.Itoa(len(bs)))
@@ -418,7 +526,7 @@ func (req *ProxyRequest) SetURLParameter(key string, value string) {
 	req.ParseForm()
 }
 
-func (req *ProxyRequest) URLParameters() (url.Values) {
+func (req *ProxyRequest) URLParameters() url.Values {
 	vals := req.URL.Query()
 	return vals
 }
@@ -479,7 +587,7 @@ func (req *ProxyRequest) StatusLine() string {
 	return fmt.Sprintf("%s %s %s", req.Method, req.HTTPPath(), req.Proto)
 }
 
-func (req *ProxyRequest) HeaderSection() (string) {
+func (req *ProxyRequest) HeaderSection() string {
 	retStr := req.StatusLine()
 	retStr += "\r\n"
 	for k, vs := range req.Header {
@@ -495,9 +603,9 @@ func (rsp *ProxyResponse) resetBodyReader() {
 	rsp.Body = ioutil.NopCloser(bytes.NewBuffer(rsp.BodyBytes()))
 }
 
-func (rsp *ProxyResponse) RepeatableWrite(w io.Writer) {
-	rsp.Write(w)
-	rsp.resetBodyReader()
+func (rsp *ProxyResponse) RepeatableWrite(w io.Writer) error {
+	defer rsp.resetBodyReader()
+	return rsp.Write(w)
 }
 
 func (rsp *ProxyResponse) BodyBytes() []byte {
@@ -510,7 +618,7 @@ func (rsp *ProxyResponse) SetBodyBytes(bs []byte) {
 	rsp.Header.Set("Content-Length", strconv.Itoa(len(bs)))
 }
 
-func (rsp *ProxyResponse) Clone() (*ProxyResponse) {
+func (rsp *ProxyResponse) Clone() *ProxyResponse {
 	buf := bytes.NewBuffer(make([]byte, 0))
 	rsp.RepeatableWrite(buf)
 	newRsp, err := ProxyResponseFromBytes(buf.Bytes())
@@ -520,7 +628,7 @@ func (rsp *ProxyResponse) Clone() (*ProxyResponse) {
 	return newRsp
 }
 
-func (rsp *ProxyResponse) DeepClone() (*ProxyResponse) {
+func (rsp *ProxyResponse) DeepClone() *ProxyResponse {
 	newRsp := rsp.Clone()
 	newRsp.DbId = rsp.DbId
 	if rsp.Unmangled != nil {
@@ -565,7 +673,7 @@ func (rsp *ProxyResponse) StatusLine() string {
 	return fmt.Sprintf("HTTP/%d.%d %03d %s", rsp.ProtoMajor, rsp.ProtoMinor, rsp.StatusCode, rsp.HTTPStatus())
 }
 
-func (rsp *ProxyResponse) HeaderSection() (string) {
+func (rsp *ProxyResponse) HeaderSection() string {
 	retStr := rsp.StatusLine()
 	retStr += "\r\n"
 	for k, vs := range rsp.Header {
@@ -585,7 +693,7 @@ func (msg *ProxyWSMessage) String() string {
 	return fmt.Sprintf("{WS Message  msg=\"%s\", type=%d, dir=%s}", string(msg.Message), msg.Type, dirStr)
 }
 
-func (msg *ProxyWSMessage) Clone() (*ProxyWSMessage) {
+func (msg *ProxyWSMessage) Clone() *ProxyWSMessage {
 	var retMsg ProxyWSMessage
 	retMsg.Type = msg.Type
 	retMsg.Message = msg.Message
@@ -595,7 +703,7 @@ func (msg *ProxyWSMessage) Clone() (*ProxyWSMessage) {
 	return &retMsg
 }
 
-func (msg *ProxyWSMessage) DeepClone() (*ProxyWSMessage) {
+func (msg *ProxyWSMessage) DeepClone() *ProxyWSMessage {
 	retMsg := msg.Clone()
 	retMsg.DbId = msg.DbId
 	if msg.Unmangled != nil {
@@ -613,7 +721,7 @@ func (msg *ProxyWSMessage) Eq(other *ProxyWSMessage) bool {
 	return true
 }
 
-func CopyHeader(hd http.Header) (http.Header) {
+func CopyHeader(hd http.Header) http.Header {
 	var ret http.Header = make(http.Header)
 	for k, vs := range hd {
 		for _, v := range vs {
@@ -622,3 +730,80 @@ func CopyHeader(hd http.Header) (http.Header) {
 	}
 	return ret
 }
+
+func submitRequest(req *ProxyRequest, useProxy bool, proxyHost string,
+	proxyPort int, proxyCreds *ProxyCredentials, proxyIsSOCKS bool) error {
+	var dialer NetDialer = req.NetDial
+	if dialer == nil {
+		dialer = net.Dial
+	}
+
+	var conn net.Conn
+	var err error
+	var proxyFormat bool = false
+	if useProxy {
+		if proxyIsSOCKS {
+			var socksCreds *proxy.Auth
+			if proxyCreds != nil {
+				socksCreds = &proxy.Auth{
+					User:     proxyCreds.Username,
+					Password: proxyCreds.Password,
+				}
+			}
+			socksDialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("%s:%d", proxyHost, proxyPort), socksCreds, proxy.Direct)
+			if err != nil {
+				return fmt.Errorf("error creating SOCKS dialer: %s", err.Error())
+			}
+			conn, err = socksDialer.Dial("tcp", fmt.Sprintf("%s:%d", req.DestHost, req.DestPort))
+			if err != nil {
+				return fmt.Errorf("error dialing host: %s", err.Error())
+			}
+			defer conn.Close()
+		} else {
+			conn, err = dialer("tcp", fmt.Sprintf("%s:%d", proxyHost, proxyPort))
+			if err != nil {
+				return fmt.Errorf("error dialing proxy: %s", err.Error())
+			}
+			defer conn.Close()
+			if req.DestUseTLS {
+				if err := PerformConnect(conn, req.DestHost, req.DestPort); err != nil {
+					return err
+				}
+				proxyFormat = false
+			} else {
+				proxyFormat = true
+			}
+		}
+	} else {
+		conn, err = dialer("tcp", fmt.Sprintf("%s:%d", req.DestHost, req.DestPort))
+		if err != nil {
+			return fmt.Errorf("error dialing host: %s", err.Error())
+		}
+		defer conn.Close()
+	}
+
+	if req.DestUseTLS {
+		tls_conn := tls.Client(conn, &tls.Config{
+			InsecureSkipVerify: true,
+		})
+		conn = tls_conn
+	}
+
+	if proxyFormat {
+		return req.SubmitProxy(conn, proxyCreds)
+	} else {
+		return req.Submit(conn)
+	}
+}
+
+func SubmitRequest(req *ProxyRequest) error {
+	return submitRequest(req, false, "", 0, nil, false)
+}
+
+func SubmitRequestProxy(req *ProxyRequest, proxyHost string, proxyPort int, creds *ProxyCredentials) error {
+	return submitRequest(req, true, proxyHost, proxyPort, creds, false)
+}
+
+func SubmitRequestSOCKSProxy(req *ProxyRequest, proxyHost string, proxyPort int, creds *ProxyCredentials) error {
+	return submitRequest(req, true, proxyHost, proxyPort, creds, true)
+}

proxyhttp_test.go

@@ -4,7 +4,6 @@ import (
 	"net/url"
 	"runtime"
 	"testing"
-
 	// "bytes"
 	// "net/http"
 	// "bufio"
@@ -80,13 +79,13 @@ func TestEq(t *testing.T) {
 	}
 
 	req1.Header = map[string][]string{
-		"Foo": []string{"Bar", "Baz"},
-		"Foo2": []string{"Bar2", "Baz2"},
+		"Foo":    []string{"Bar", "Baz"},
+		"Foo2":   []string{"Bar2", "Baz2"},
 		"Cookie": []string{"cookie=cocks"},
 	}
 	req2.Header = map[string][]string{
-		"Foo": []string{"Bar", "Baz"},
-		"Foo2": []string{"Bar2", "Baz2"},
+		"Foo":    []string{"Bar", "Baz"},
+		"Foo2":   []string{"Bar2", "Baz2"},
 		"Cookie": []string{"cookie=cocks"},
 	}
 
@@ -95,8 +94,8 @@ func TestEq(t *testing.T) {
 	}
 
 	req2.Header = map[string][]string{
-		"Foo": []string{"Baz", "Bar"},
-		"Foo2": []string{"Bar2", "Baz2"},
+		"Foo":    []string{"Baz", "Bar"},
+		"Foo2":   []string{"Bar2", "Baz2"},
 		"Cookie": []string{"cookie=cocks"},
 	}
 	if req1.Eq(req2) {
@@ -104,8 +103,8 @@ func TestEq(t *testing.T) {
 	}
 
 	req2.Header = map[string][]string{
-		"Foo": []string{"Bar", "Baz"},
-		"Foo2": []string{"Bar2", "Baz2"},
+		"Foo":    []string{"Bar", "Baz"},
+		"Foo2":   []string{"Bar2", "Baz2"},
 		"Cookie": []string{"cookiee=cocks"},
 	}
 	if req1.Eq(req2) {

proxylistener.go

@@ -54,7 +54,7 @@ type ProxyConn interface {
 	net.Conn
 
 	Id() int
-	Logger() (*log.Logger)
+	Logger() *log.Logger
 
 	SetCACertificate(*tls.Certificate)
 	StartMaybeTLS(hostname string) (bool, error)
@@ -67,12 +67,12 @@ type proxyAddr struct {
 }
 
 type proxyConn struct {
-	Addr   *proxyAddr
-	logger *log.Logger
-	id     int
+	Addr    *proxyAddr
+	logger  *log.Logger
+	id      int
 	conn    net.Conn      // Wrapped connection
 	readReq *http.Request // A replaced request
-	caCert *tls.Certificate
+	caCert  *tls.Certificate
 }
 
 // ProxyAddr implementations/functions
@@ -123,7 +123,6 @@ func (a *proxyAddr) String() string {
 	return EncodeRemoteAddr(a.Host, a.Port, a.UseTLS)
 }
 
-
 //// bufferedConn and wrappers
 type bufferedConn struct {
 	reader   *bufio.Reader
@@ -227,7 +226,7 @@ func (pconn *proxyConn) StartMaybeTLS(hostname string) (bool, error) {
 
 		config := &tls.Config{
 			InsecureSkipVerify: true,
-			Certificates: []tls.Certificate{cert},
+			Certificates:       []tls.Certificate{cert},
 		}
 		tlsConn := tls.Server(bufConn, config)
 		pconn.conn = tlsConn
@@ -239,8 +238,8 @@ func (pconn *proxyConn) StartMaybeTLS(hostname string) (bool, error) {
 }
 
 func NewProxyConn(c net.Conn, l *log.Logger) *proxyConn {
-	a := proxyAddr{Host:"", Port:-1, UseTLS:false}
-	p := proxyConn{Addr:&a, logger:l, conn:c, readReq:nil}
+	a := proxyAddr{Host: "", Port: -1, UseTLS: false}
+	p := proxyConn{Addr: &a, logger: l, conn: c, readReq: nil}
 	p.id = getNextConnId()
 	return &p
 }
@@ -262,15 +261,15 @@ type ProxyListener struct {
 
 	State int
 
-	inputListeners   mapset.Set
-	mtx              sync.Mutex
-	logger           *log.Logger
-	outputConns      chan ProxyConn
-	inputConns       chan net.Conn
-	outputConnDone   chan struct{}
-	inputConnDone    chan struct{}
-	listenWg         sync.WaitGroup
-	caCert           *tls.Certificate
+	inputListeners mapset.Set
+	mtx            sync.Mutex
+	logger         *log.Logger
+	outputConns    chan ProxyConn
+	inputConns     chan net.Conn
+	outputConnDone chan struct{}
+	inputConnDone  chan struct{}
+	listenWg       sync.WaitGroup
+	caCert         *tls.Certificate
 }
 
 type listenerData struct {
@@ -278,7 +277,7 @@ type listenerData struct {
 	Listener net.Listener
 }
 
-func newListenerData(listener net.Listener) (*listenerData) {
+func newListenerData(listener net.Listener) *listenerData {
 	l := listenerData{}
 	l.Id = getNextListenerId()
 	l.Listener = listener
@@ -372,7 +371,7 @@ func (listener *ProxyListener) AddListener(inlisten net.Listener) error {
 
 	listener.logger.Println("Adding listener to ProxyListener:", inlisten)
 	il := newListenerData(inlisten)
-	l  := listener
+	l := listener
 	listener.listenWg.Add(1)
 	go func() {
 		defer l.listenWg.Done()

proxymessages.go

@@ -1,9 +1,9 @@
 package main
 
 import (
-	"crypto/tls"
 	"bufio"
 	"bytes"
+	"crypto/tls"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -52,53 +52,54 @@ func NewProxyMessageListener(logger *log.Logger, iproxy *InterceptingProxy) *Mes
 	l.AddHandler("closestorage", closeStorageHandler)
 	l.AddHandler("setproxystorage", setProxyStorageHandler)
 	l.AddHandler("liststorage", listProxyStorageHandler)
+	l.AddHandler("setproxy", setProxyHandler)
 
 	return l
 }
 
 // Message input structs
 type RequestJSON struct {
-	DestHost string
-	DestPort int
-	UseTLS bool
-	Method string
-	Path string
+	DestHost   string
+	DestPort   int
+	UseTLS     bool
+	Method     string
+	Path       string
 	ProtoMajor int
 	ProtoMinor int
-	Headers map[string][]string
-	Body string
-	Tags []string
+	Headers    map[string][]string
+	Body       string
+	Tags       []string
 
 	StartTime int64 `json:"StartTime,omitempty"`
-	EndTime int64 `json:"EndTime,omitempty"`
+	EndTime   int64 `json:"EndTime,omitempty"`
 
-	Unmangled *RequestJSON `json:"Unmangled,omitempty"`
-	Response *ResponseJSON `json:"Response,omitempty"`
+	Unmangled  *RequestJSON     `json:"Unmangled,omitempty"`
+	Response   *ResponseJSON    `json:"Response,omitempty"`
 	WSMessages []*WSMessageJSON `json:"WSMessages,omitempty"`
-	DbId string `json:"DbId,omitempty"`
+	DbId       string           `json:"DbId,omitempty"`
 }
 
 type ResponseJSON struct {
 	ProtoMajor int
 	ProtoMinor int
 	StatusCode int
-	Reason string
+	Reason     string
 
 	Headers map[string][]string
-	Body string
+	Body    string
 
 	Unmangled *ResponseJSON `json:"Unmangled,omitempty"`
-	DbId string
+	DbId      string
 }
 
 type WSMessageJSON struct {
-	Message string
-	IsBinary bool
-	ToServer bool
+	Message   string
+	IsBinary  bool
+	ToServer  bool
 	Timestamp int64
 
 	Unmangled *WSMessageJSON `json:"Unmangled,omitempty"`
-	DbId string
+	DbId      string
 }
 
 func (reqd *RequestJSON) Validate() error {
@@ -162,7 +163,7 @@ func (reqd *RequestJSON) Parse() (*ProxyRequest, error) {
 		req.EndDatetime = time.Unix(0, reqd.EndTime)
 	}
 
-	for _, tag := range(reqd.Tags) {
+	for _, tag := range reqd.Tags {
 		req.AddTag(tag)
 	}
 
@@ -186,7 +187,7 @@ func (reqd *RequestJSON) Parse() (*ProxyRequest, error) {
 	return req, nil
 }
 
-func NewRequestJSON(req *ProxyRequest, headersOnly bool) (*RequestJSON) {
+func NewRequestJSON(req *ProxyRequest, headersOnly bool) *RequestJSON {
 
 	newHeaders := make(map[string][]string)
 	for k, vs := range req.Header {
@@ -217,23 +218,23 @@ func NewRequestJSON(req *ProxyRequest, headersOnly bool) (*RequestJSON) {
 	}
 
 	ret := &RequestJSON{
-		DestHost: req.DestHost,
-		DestPort: req.DestPort,
-		UseTLS: req.DestUseTLS,
-		Method: req.Method,
-		Path: req.HTTPPath(),
+		DestHost:   req.DestHost,
+		DestPort:   req.DestPort,
+		UseTLS:     req.DestUseTLS,
+		Method:     req.Method,
+		Path:       req.HTTPPath(),
 		ProtoMajor: req.ProtoMajor,
 		ProtoMinor: req.ProtoMinor,
-		Headers: newHeaders,
-		Tags: req.Tags(),
+		Headers:    newHeaders,
+		Tags:       req.Tags(),
 
 		StartTime: req.StartDatetime.UnixNano(),
-		EndTime: req.EndDatetime.UnixNano(),
+		EndTime:   req.EndDatetime.UnixNano(),
 
-		Unmangled: unmangled,
-		Response: rsp,
+		Unmangled:  unmangled,
+		Response:   rsp,
 		WSMessages: wsms,
-		DbId: req.DbId,
+		DbId:       req.DbId,
 	}
 	if !headersOnly {
 		ret.Body = base64.StdEncoding.EncodeToString(req.BodyBytes())
@@ -294,7 +295,7 @@ func (rspd *ResponseJSON) Parse() (*ProxyResponse, error) {
 	return rsp, nil
 }
 
-func NewResponseJSON(rsp *ProxyResponse, headersOnly bool) (*ResponseJSON) {
+func NewResponseJSON(rsp *ProxyResponse, headersOnly bool) *ResponseJSON {
 	newHeaders := make(map[string][]string)
 	for k, vs := range rsp.Header {
 		for _, v := range vs {
@@ -317,10 +318,10 @@ func NewResponseJSON(rsp *ProxyResponse, headersOnly bool) (*ResponseJSON) {
 		ProtoMajor: rsp.ProtoMajor,
 		ProtoMinor: rsp.ProtoMinor,
 		StatusCode: rsp.StatusCode,
-		Reason: rsp.HTTPStatus(),
-		Headers: newHeaders,
-		DbId: rsp.DbId,
-		Unmangled: unmangled,
+		Reason:     rsp.HTTPStatus(),
+		Headers:    newHeaders,
+		DbId:       rsp.DbId,
+		Unmangled:  unmangled,
 	}
 
 	if !headersOnly {
@@ -359,18 +360,18 @@ func (wsmd *WSMessageJSON) Parse() (*ProxyWSMessage, error) {
 	}
 
 	retData := &ProxyWSMessage{
-		Message: message,
-		Type: mtype,
+		Message:   message,
+		Type:      mtype,
 		Direction: Direction,
 		Timestamp: time.Unix(0, wsmd.Timestamp),
 		Unmangled: unmangled,
-		DbId: wsmd.DbId,
+		DbId:      wsmd.DbId,
 	}
 
 	return retData, nil
 }
 
-func NewWSMessageJSON(wsm *ProxyWSMessage) (*WSMessageJSON) {
+func NewWSMessageJSON(wsm *ProxyWSMessage) *WSMessageJSON {
 	isBinary := false
 	if wsm.Type == websocket.BinaryMessage {
 		isBinary = true
@@ -387,12 +388,12 @@ func NewWSMessageJSON(wsm *ProxyWSMessage) (*WSMessageJSON) {
 	}
 
 	ret := &WSMessageJSON{
-		Message: base64.StdEncoding.EncodeToString(wsm.Message),
-		IsBinary: isBinary,
-		ToServer: toServer,
+		Message:   base64.StdEncoding.EncodeToString(wsm.Message),
+		IsBinary:  isBinary,
+		ToServer:  toServer,
 		Timestamp: wsm.Timestamp.UnixNano(),
 		Unmangled: unmangled,
-		DbId: wsm.DbId,
+		DbId:      wsm.DbId,
 	}
 
 	return ret
@@ -427,11 +428,11 @@ type successResult struct {
 /*
 Ping
 */
-type pingMessage struct {}
+type pingMessage struct{}
 
 type pingResponse struct {
 	Success bool
-	Ping string
+	Ping    string
 }
 
 func pingHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
@@ -448,7 +449,7 @@ type submitMessage struct {
 }
 
 type submitResponse struct {
-	Success bool
+	Success          bool
 	SubmittedRequest *RequestJSON
 }
 
@@ -475,8 +476,8 @@ func submitHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Interceptin
 		}
 		SaveNewRequest(storage, req)
 	}
-	logger.Println("Submitting request to", req.FullURL(),"...")
-	if err := req.Submit(); err != nil {
+	logger.Println("Submitting request to", req.FullURL(), "...")
+	if err := iproxy.SubmitRequest(req); err != nil {
 		ErrorResponse(c, err.Error())
 		return
 	}
@@ -504,7 +505,7 @@ type saveNewMessage struct {
 
 type saveNewResponse struct {
 	Success bool
-	DbId string
+	DbId    string
 }
 
 func saveNewHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
@@ -540,7 +541,7 @@ func saveNewHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercepti
 
 	response := &saveNewResponse{
 		Success: true,
-		DbId: req.DbId,
+		DbId:    req.DbId,
 	}
 	MessageResponse(c, response)
 }
@@ -549,10 +550,10 @@ func saveNewHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercepti
 QueryRequests
 */
 type storageQueryMessage struct {
-	Query StrMessageQuery
+	Query       StrMessageQuery
 	HeadersOnly bool
-	MaxResults int64
-	Storage int
+	MaxResults  int64
+	Storage     int
 }
 
 type storageQueryResult struct {
@@ -562,9 +563,9 @@ type storageQueryResult struct {
 
 func storageQueryHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
 	mreq := storageQueryMessage{
-		Query: nil,
+		Query:       nil,
 		HeadersOnly: false,
-		MaxResults: 0,
+		MaxResults:  0,
 	}
 
 	if err := json.Unmarshal(b, &mreq); err != nil {
@@ -658,10 +659,9 @@ func validateQueryHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Inte
 		ErrorResponse(c, err.Error())
 		return
 	}
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
-
 /*
 SetScope
 */
@@ -690,7 +690,7 @@ func setScopeHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercept
 		return
 	}
 
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 /*
@@ -701,9 +701,9 @@ type viewScopeMessage struct {
 }
 
 type viewScopeResult struct {
-	Success bool
+	Success  bool
 	IsCustom bool
-	Query StrMessageQuery
+	Query    StrMessageQuery
 }
 
 func viewScopeHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
@@ -712,7 +712,7 @@ func viewScopeHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 
 	if scopeQuery == nil && scopeChecker != nil {
 		MessageResponse(c, &viewScopeResult{
-			Success: true,
+			Success:  true,
 			IsCustom: true,
 		})
 		return
@@ -726,9 +726,9 @@ func viewScopeHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 	}
 
 	MessageResponse(c, &viewScopeResult{
-		Success: true,
+		Success:  true,
 		IsCustom: false,
-		Query: strQuery,
+		Query:    strQuery,
 	})
 }
 
@@ -737,8 +737,8 @@ Tag messages
 */
 
 type addTagMessage struct {
-	ReqId string
-	Tag string
+	ReqId   string
+	Tag     string
 	Storage int
 }
 
@@ -779,12 +779,12 @@ func addTagHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Interceptin
 		return
 	}
 
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type removeTagMessage struct {
-	ReqId string
-	Tag string
+	ReqId   string
+	Tag     string
 	Storage int
 }
 
@@ -825,11 +825,11 @@ func removeTagHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 		return
 	}
 
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type clearTagsMessage struct {
-	ReqId string
+	ReqId   string
 	Storage int
 }
 
@@ -866,7 +866,7 @@ func clearTagHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercept
 		return
 	}
 
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 /*
@@ -874,12 +874,12 @@ Intercept
 */
 
 type interceptMessage struct {
-	InterceptRequests bool
+	InterceptRequests  bool
 	InterceptResponses bool
-	InterceptWS bool
+	InterceptWS        bool
 
 	UseQuery bool
-	Query MessageQuery
+	Query    MessageQuery
 }
 
 type intHandshakeResult struct {
@@ -891,49 +891,49 @@ var getNextIntId = IdCounter()
 
 type intRequest struct {
 	// A request to have a message mangled
-	Id int
-	Type string
+	Id      int
+	Type    string
 	Success bool
-	Result chan *intResponse `json:"-"`
+	Result  chan *intResponse `json:"-"`
 
-	Request *RequestJSON `json:"Request,omitempty"`
-	Response *ResponseJSON `json:"Response,omitempty"`
+	Request   *RequestJSON   `json:"Request,omitempty"`
+	Response  *ResponseJSON  `json:"Response,omitempty"`
 	WSMessage *WSMessageJSON `json:"WSMessage,omitempty"`
 }
 
 type intResponse struct {
 	// response from the client with a mangled http request
-	Id int
+	Id      int
 	Dropped bool
 
-	Request *RequestJSON `json:"Request,omitempty"`
-	Response *ResponseJSON `json:"Response,omitempty"`
+	Request   *RequestJSON   `json:"Request,omitempty"`
+	Response  *ResponseJSON  `json:"Response,omitempty"`
 	WSMessage *WSMessageJSON `json:"WSMessage,omitempty"`
 }
 
 type intErrorMessage struct {
 	// a message template for sending an error to client if there is an error
 	// with the mangled message they sent
-	Id int
+	Id      int
 	Success bool
-	Reason string
+	Reason  string
 }
 
 func intErrorResponse(id int, conn net.Conn, reason string) {
 	m := &intErrorMessage{
-		Id: id,
+		Id:      id,
 		Success: false,
-		Reason: reason,
+		Reason:  reason,
 	}
 	MessageResponse(conn, m)
 }
 
 func interceptHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
 	mreq := interceptMessage{
-		InterceptRequests: false,
+		InterceptRequests:  false,
 		InterceptResponses: false,
-		InterceptWS: false,
-		UseQuery: false,
+		InterceptWS:        false,
+		UseQuery:           false,
 	}
 
 	if err := json.Unmarshal(b, &mreq); err != nil {
@@ -1002,9 +1002,9 @@ func interceptHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 
 			// convert request data to an intRequest
 			intReq := &intRequest{
-				Id: getNextIntId(),
-				Type: "httprequest",
-				Result: make(chan *intResponse),
+				Id:      getNextIntId(),
+				Type:    "httprequest",
+				Result:  make(chan *intResponse),
 				Success: true,
 
 				Request: reqData,
@@ -1059,12 +1059,12 @@ func interceptHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 			CleanRspJSON(rspData)
 
 			intReq := &intRequest{
-				Id: getNextIntId(),
-				Type: "httpresponse",
-				Result: make(chan *intResponse),
+				Id:      getNextIntId(),
+				Type:    "httpresponse",
+				Result:  make(chan *intResponse),
 				Success: true,
 
-				Request: reqData,
+				Request:  reqData,
 				Response: rspData,
 			}
 
@@ -1124,15 +1124,14 @@ func interceptHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 			rspData := NewResponseJSON(rsp, false)
 			CleanRspJSON(rspData)
 
-
 			intReq := &intRequest{
-				Id: getNextIntId(),
-				Type: msgType,
-				Result: make(chan *intResponse),
+				Id:      getNextIntId(),
+				Type:    msgType,
+				Result:  make(chan *intResponse),
 				Success: true,
 
-				Request: reqData,
-				Response: rspData,
+				Request:   reqData,
+				Response:  rspData,
 				WSMessage: wsData,
 			}
 
@@ -1296,7 +1295,7 @@ func allSavedQueriesHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *In
 	savedQueries := make([]*StrSavedQuery, 0)
 	for _, q := range goQueries {
 		strSavedQuery := &StrSavedQuery{
-			Name: q.Name,
+			Name:  q.Name,
 			Query: nil,
 		}
 		sq, err := GoQueryToStrQuery(q.Query)
@@ -1312,8 +1311,8 @@ func allSavedQueriesHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *In
 }
 
 type saveQueryMessage struct {
-	Name string
-	Query StrMessageQuery
+	Name    string
+	Query   StrMessageQuery
 	Storage int
 }
 
@@ -1358,17 +1357,17 @@ func saveQueryHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 		return
 	}
 
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type loadQueryMessage struct {
-	Name string
+	Name    string
 	Storage int
 }
 
 type loadQueryResult struct {
 	Success bool
-	Query StrMessageQuery
+	Query   StrMessageQuery
 }
 
 func loadQueryHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
@@ -1408,14 +1407,14 @@ func loadQueryHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Intercep
 
 	result := &loadQueryResult{
 		Success: true,
-		Query: strQuery,
+		Query:   strQuery,
 	}
 
 	MessageResponse(c, result)
 }
 
 type deleteQueryMessage struct {
-	Name string
+	Name    string
 	Storage int
 }
 
@@ -1442,7 +1441,7 @@ func deleteQueryHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Interc
 		ErrorResponse(c, err.Error())
 		return
 	}
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 /*
@@ -1450,10 +1449,10 @@ Listener management
 */
 
 type activeListener struct {
-	Id int
+	Id       int
 	Listener net.Listener `json:"-"`
-	Type string
-	Addr string
+	Type     string
+	Addr     string
 }
 
 type addListenerMessage struct {
@@ -1463,7 +1462,7 @@ type addListenerMessage struct {
 
 type addListenerResult struct {
 	Success bool
-	Id int
+	Id      int
 }
 
 var getNextMsgListenerId = IdCounter()
@@ -1500,10 +1499,10 @@ func addListenerHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Interc
 	iproxy.AddListener(listener)
 
 	alistener := &activeListener{
-		Id: getNextMsgListenerId(),
+		Id:       getNextMsgListenerId(),
 		Listener: listener,
-		Type: mreq.Type,
-		Addr: mreq.Addr,
+		Type:     mreq.Type,
+		Addr:     mreq.Addr,
 	}
 
 	msgActiveListenersMtx.Lock()
@@ -1511,7 +1510,7 @@ func addListenerHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Interc
 	msgActiveListeners[alistener.Id] = alistener
 	result := &addListenerResult{
 		Success: true,
-		Id: alistener.Id,
+		Id:      alistener.Id,
 	}
 
 	MessageResponse(c, result)
@@ -1538,11 +1537,10 @@ func removeListenerHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Int
 
 	iproxy.RemoveListener(alistener.Listener)
 	delete(msgActiveListeners, alistener.Id)
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
-
-type getListenersMessage struct {}
+type getListenersMessage struct{}
 
 type getListenersResult struct {
 	Success bool
@@ -1567,7 +1565,7 @@ Certificate Management
 */
 
 type loadCertificatesMessage struct {
-	KeyFile string
+	KeyFile         string
 	CertificateFile string
 }
 
@@ -1583,18 +1581,18 @@ func loadCertificatesHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *I
 		return
 	}
 
- 	caCert, err := tls.LoadX509KeyPair(mreq.CertificateFile, mreq.KeyFile)
+	caCert, err := tls.LoadX509KeyPair(mreq.CertificateFile, mreq.KeyFile)
 	if err != nil {
 		ErrorResponse(c, err.Error())
 		return
 	}
 
 	iproxy.SetCACertificate(&caCert)
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type setCertificatesMessage struct {
-	KeyPEMData []byte
+	KeyPEMData         []byte
 	CertificatePEMData []byte
 }
 
@@ -1610,23 +1608,23 @@ func setCertificatesHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *In
 		return
 	}
 
- 	caCert, err := tls.X509KeyPair(mreq.CertificatePEMData, mreq.KeyPEMData)
+	caCert, err := tls.X509KeyPair(mreq.CertificatePEMData, mreq.KeyPEMData)
 	if err != nil {
 		ErrorResponse(c, err.Error())
 		return
 	}
 
 	iproxy.SetCACertificate(&caCert)
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 func clearCertificatesHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
 	iproxy.SetCACertificate(nil)
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type generateCertificatesMessage struct {
-	KeyFile string
+	KeyFile  string
 	CertFile string
 }
 
@@ -1639,13 +1637,13 @@ func generateCertificatesHandler(b []byte, c net.Conn, logger *log.Logger, iprox
 
 	pair, err := GenerateCACerts()
 	if err != nil {
-		ErrorResponse(c, "error generating certificates: " + err.Error())
+		ErrorResponse(c, "error generating certificates: "+err.Error())
 		return
 	}
 
 	pkeyFile, err := os.OpenFile(mreq.KeyFile, os.O_RDWR|os.O_CREATE, 0600)
 	if err != nil {
-		ErrorResponse(c, "could not save private key: " + err.Error())
+		ErrorResponse(c, "could not save private key: "+err.Error())
 		return
 	}
 	pkeyFile.Write(pair.PrivateKeyPEM())
@@ -1656,7 +1654,7 @@ func generateCertificatesHandler(b []byte, c net.Conn, logger *log.Logger, iprox
 
 	certFile, err := os.OpenFile(mreq.CertFile, os.O_RDWR|os.O_CREATE, 0600)
 	if err != nil {
-		ErrorResponse(c, "could not save private key: " + err.Error())
+		ErrorResponse(c, "could not save private key: "+err.Error())
 		return
 	}
 	certFile.Write(pair.CACertPEM())
@@ -1665,12 +1663,12 @@ func generateCertificatesHandler(b []byte, c net.Conn, logger *log.Logger, iprox
 		return
 	}
 
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type generatePEMCertificatesResult struct {
-	Success bool
-	KeyPEMData []byte
+	Success            bool
+	KeyPEMData         []byte
 	CertificatePEMData []byte
 }
 
@@ -1683,13 +1681,13 @@ func generatePEMCertificatesHandler(b []byte, c net.Conn, logger *log.Logger, ip
 
 	pair, err := GenerateCACerts()
 	if err != nil {
-		ErrorResponse(c, "error generating certificates: " + err.Error())
+		ErrorResponse(c, "error generating certificates: "+err.Error())
 		return
 	}
 
 	result := &generatePEMCertificatesResult{
-		Success: true,
-		KeyPEMData: pair.PrivateKeyPEM(),
+		Success:            true,
+		KeyPEMData:         pair.PrivateKeyPEM(),
 		CertificatePEMData: pair.CACertPEM(),
 	}
 	MessageResponse(c, result)
@@ -1700,12 +1698,12 @@ Storage functions
 */
 
 type addSQLiteStorageMessage struct {
-	Path string
+	Path        string
 	Description string
 }
 
 type addSQLiteStorageResult struct {
-	Success bool
+	Success   bool
 	StorageId int
 }
 
@@ -1723,13 +1721,13 @@ func addSQLiteStorageHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *I
 
 	storage, err := OpenSQLiteStorage(mreq.Path, logger)
 	if err != nil {
-		ErrorResponse(c, "error opening SQLite databae: " + err.Error())
+		ErrorResponse(c, "error opening SQLite databae: "+err.Error())
 		return
 	}
 
 	sid := iproxy.AddMessageStorage(storage, mreq.Description)
 	result := &addSQLiteStorageResult{
-		Success: true,
+		Success:   true,
 		StorageId: sid,
 	}
 	MessageResponse(c, result)
@@ -1740,7 +1738,7 @@ type addInMemoryStorageMessage struct {
 }
 
 type addInMemoryStorageResult struct {
-	Success bool
+	Success   bool
 	StorageId int
 }
 
@@ -1753,13 +1751,13 @@ func addInMemoryStorageHandler(b []byte, c net.Conn, logger *log.Logger, iproxy
 
 	storage, err := InMemoryStorage(logger)
 	if err != nil {
-		ErrorResponse(c, "error creating in memory storage: " + err.Error())
+		ErrorResponse(c, "error creating in memory storage: "+err.Error())
 		return
 	}
 
 	sid := iproxy.AddMessageStorage(storage, mreq.Description)
 	result := &addInMemoryStorageResult{
-		Success: true,
+		Success:   true,
 		StorageId: sid,
 	}
 	MessageResponse(c, result)
@@ -1786,7 +1784,7 @@ func closeStorageHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *Inter
 	}
 
 	iproxy.CloseMessageStorage(mreq.StorageId)
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type setProxyStorageMessage struct {
@@ -1811,17 +1809,17 @@ func setProxyStorageHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *In
 		return
 	}
 
-	MessageResponse(c, &successResult{Success:true})
+	MessageResponse(c, &successResult{Success: true})
 }
 
 type savedStorageJSON struct {
-	Id int
+	Id          int
 	Description string
 }
 
 type listProxyStorageResult struct {
 	Storages []*savedStorageJSON
-	Success bool
+	Success  bool
 }
 
 func listProxyStorageHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
@@ -1832,7 +1830,49 @@ func listProxyStorageHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *I
 	}
 	result := &listProxyStorageResult{
 		Storages: storagesJSON,
-		Success: true,
+		Success:  true,
 	}
 	MessageResponse(c, result)
 }
+
+/*
+SetProxy
+*/
+
+type setProxyMessage struct {
+	UseProxy       bool
+	ProxyHost      string
+	ProxyPort      int
+	ProxyIsSOCKS   bool
+	UseCredentials bool
+	Username       string
+	Password       string
+}
+
+func setProxyHandler(b []byte, c net.Conn, logger *log.Logger, iproxy *InterceptingProxy) {
+	mreq := setProxyMessage{}
+	if err := json.Unmarshal(b, &mreq); err != nil {
+		ErrorResponse(c, "error parsing message")
+		return
+	}
+
+	var creds *ProxyCredentials = nil
+	if mreq.UseCredentials {
+		creds = &ProxyCredentials{
+			Username: mreq.Username,
+			Password: mreq.Password,
+		}
+	}
+
+	if !mreq.UseProxy {
+		iproxy.ClearUpstreamProxy()
+	} else {
+		if mreq.ProxyIsSOCKS {
+			iproxy.SetUpstreamSOCKSProxy(mreq.ProxyHost, mreq.ProxyPort, creds)
+		} else {
+			iproxy.SetUpstreamProxy(mreq.ProxyHost, mreq.ProxyPort, creds)
+		}
+	}
+
+	MessageResponse(c, &successResult{Success: true})
+}

python/puppy/puppyproxy/config.py

@@ -4,7 +4,8 @@ import json
 default_config = """{
     "listeners": [
         {"iface": "127.0.0.1", "port": 8080}
-    ]
+    ],
+    "proxy": {"use_proxy": false, "host": "", "port": 0, "is_socks": false}
 }"""
 
 
@@ -12,6 +13,7 @@ class ProxyConfig:
     
     def __init__(self):
         self._listeners = [('127.0.0.1', '8080')]
+        self._proxy = {'use_proxy': False, 'host': '', 'port': 0, 'is_socks': False}
         
     def load(self, fname):
         try:
@@ -40,6 +42,10 @@ class ProxyConfig:
 
                 self._listeners.append((iface, port))
 
+        if 'proxy' in config_info:
+            self._proxy = config_info['proxy']
+            
+                
     @property
     def listeners(self):
         return copy.deepcopy(self._listeners)
@@ -47,3 +53,67 @@ class ProxyConfig:
     @listeners.setter
     def listeners(self, val):
         self._listeners = val
+        
+    @property
+    def proxy(self):
+        # don't use this, use the getters to get the parsed values
+        return self._proxy
+        
+    @proxy.setter
+    def proxy(self, val):
+        self._proxy = val
+
+    @property
+    def use_proxy(self):
+        if self._proxy is None:
+            return False
+        if 'use_proxy' in self._proxy:
+            if self._proxy['use_proxy']:
+                return True
+        return False
+
+    @property
+    def proxy_host(self):
+        if self._proxy is None:
+            return ''
+        if 'host' in self._proxy:
+            return self._proxy['host']
+        return ''
+
+    @property
+    def proxy_port(self):
+        if self._proxy is None:
+            return ''
+        if 'port' in self._proxy:
+            return self._proxy['port']
+        return ''
+
+    @property
+    def proxy_username(self):
+        if self._proxy is None:
+            return ''
+        if 'username' in self._proxy:
+            return self._proxy['username']
+        return ''
+
+    @property
+    def proxy_password(self):
+        if self._proxy is None:
+            return ''
+        if 'password' in self._proxy:
+            return self._proxy['password']
+        return ''
+
+    @property
+    def use_proxy_creds(self):
+        return ('username' in self._proxy or 'password' in self._proxy)
+
+    @property
+    def is_socks_proxy(self):
+        if self._proxy is None:
+            return False
+        if 'is_socks' in self._proxy:
+            if self._proxy['is_socks']:
+                return True
+        return False
+

python/puppy/puppyproxy/interface/context.py

@@ -1,6 +1,6 @@
 from itertools import groupby
 
-from ..proxy import InvalidQuery
+from ..proxy import InvalidQuery, time_to_nsecs
 from ..colors import Colors, Styles
 
 # class BuiltinFilters(object):
@@ -71,6 +71,11 @@ def filtercmd(client, args):
     """
     try:
         phrases = [list(group) for k, group in groupby(args, lambda x: x == "OR") if not k]
+        for phrase in phrases:
+            # we do before/after by id not by timestamp
+            if phrase[0] in ('before', 'b4', 'after', 'af') and len(phrase) > 1:
+                r = client.req_by_id(phrase[1], headers_only=True)
+                phrase[1] = str(time_to_nsecs(r.time_start))
         client.context.apply_phrase(phrases)
     except InvalidQuery as e:
         print(e)

python/puppy/puppyproxy/interface/decode.py

@@ -7,31 +7,32 @@ import string
 import urllib
 
 from ..util import hexdump, printable_data, copy_to_clipboard, clipboard_contents, encode_basic_auth, parse_basic_auth
+from ..console import CommandError
 from io import StringIO
 
 def print_maybe_bin(s):
     binary = False
     for c in s:
-        if str(c) not in string.printable:
+        if chr(c) not in string.printable:
             binary = True
             break
     if binary:
         print(hexdump(s))
     else:
-        print(s)
+        print(s.decode())
         
 def asciihex_encode_helper(s):
-    return ''.join('{0:x}'.format(c) for c in s)
+    return ''.join('{0:x}'.format(c) for c in s).encode()
 
 def asciihex_decode_helper(s):
     ret = []
     try:
         for a, b in zip(s[0::2], s[1::2]):
-            c = a+b
+            c = chr(a)+chr(b)
             ret.append(chr(int(c, 16)))
-        return ''.join(ret)
+        return ''.join(ret).encode()
     except Exception as e:
-        raise PappyException(e)
+        raise CommandError(e)
     
 def gzip_encode_helper(s):
     out = StringIO.StringIO()
@@ -54,13 +55,21 @@ def base64_decode_helper(s):
                 return s_padded
             except:
                 pass
-        raise PappyException("Unable to base64 decode string")
+        raise CommandError("Unable to base64 decode string")
+    
+def url_decode_helper(s):
+    bs = s.decode()
+    return urllib.parse.unquote(bs).encode()
+
+def url_encode_helper(s):
+    bs = s.decode()
+    return urllib.parse.quote_plus(bs).encode()
 
 def html_encode_helper(s):
-    return ''.join(['&#x{0:x};'.format(c) for c in s])
+    return ''.join(['&#x{0:x};'.format(c) for c in s]).encode()
 
 def html_decode_helper(s):
-    return html.unescape(s)
+    return html.unescape(s.decode()).encode()
 
 def _code_helper(args, func, copy=True):
     if len(args) == 0:
@@ -107,7 +116,7 @@ def url_decode(client, args):
     If no string is given, will decode the contents of the clipboard.
     Results are copied to the clipboard.
     """
-    print_maybe_bin(_code_helper(args, urllib.unquote))
+    print_maybe_bin(_code_helper(args, url_decode_helper))
 
 def url_encode(client, args):
     """
@@ -115,7 +124,7 @@ def url_encode(client, args):
     If no string is given, will encode the contents of the clipboard.
     Results are copied to the clipboard.
     """
-    print_maybe_bin(_code_helper(args, urllib.quote_plus))
+    print_maybe_bin(_code_helper(args, url_encode_helper))
 
 def asciihex_decode(client, args):
     """
@@ -187,7 +196,7 @@ def url_decode_raw(client, args):
     results will not be copied. It is suggested you redirect the output
     to a file.
     """
-    print(_code_helper(args, urllib.unquote, copy=False))
+    print(_code_helper(args, url_decode_helper, copy=False))
 
 def url_encode_raw(client, args):
     """
@@ -195,7 +204,7 @@ def url_encode_raw(client, args):
     results will not be copied. It is suggested you redirect the output
     to a file.
     """
-    print(_code_helper(args, urllib.quote_plus, copy=False))
+    print(_code_helper(args, url_encode_helper, copy=False))
 
 def asciihex_decode_raw(client, args):
     """
@@ -254,9 +263,8 @@ def unix_time_decode(client, args):
     print(_code_helper(args, unix_time_decode_helper))
     
 def http_auth_encode(client, args):
-    args = shlex.split(args[0])
     if len(args) != 2:
-        raise PappyException('Usage: http_auth_encode <username> <password>')
+        raise CommandError('Usage: http_auth_encode <username> <password>')
     username, password = args
     print(encode_basic_auth(username, password))
 

python/puppy/puppyproxy/interface/misc.py

@@ -23,7 +23,7 @@ class WatchMacro(InterceptMacro):
         printstr = "< "
         printstr += verb_color(request.method) + request.method + Colors.ENDC + ' '
         printstr += url_formatter(request, colored=True)
-        printstr += " -> "
+        printstr += " \u2192 "
         response_code = str(response.status_code) + ' ' + response.reason
         response_code = scode_color(response_code) + response_code + Colors.ENDC
         printstr += response_code

python/puppy/puppyproxy/interface/repeater/repeater.py

@@ -1524,7 +1524,7 @@ def update_buffers(req):
 
     # Save the port, ssl, host setting
     vim.command("let s:dest_port=%d" % req.dest_port)
-    vim.command("let s:dest_host='%s'" % req.dest_host)
+    vim.command("let s:dest_host='%s'" % escape(req.dest_host))
 
     if req.use_tls:
         vim.command("let s:use_tls=1")
@@ -1545,6 +1545,8 @@ def set_up_windows():
     storage_id = vim.eval("a:3")
     msg_addr = vim.eval("a:4")
 
+    vim.command("let s:storage_id=%d" % int(storage_id))
+    
     # Get the left buffer
     vim.command("new")
     vim.command("only")
@@ -1568,11 +1570,12 @@ def dest_loc():
     dest_host = vim.eval("s:dest_host")
     dest_port = int(vim.eval("s:dest_port"))
     tls_num = vim.eval("s:use_tls")
+    storage_id = int(vim.eval("s:storage_id"))
     if tls_num == "1":
         use_tls = True
     else:
         use_tls = False
-    return (dest_host, dest_port, use_tls)
+    return (dest_host, dest_port, use_tls, storage_id)
 
 def submit_current_buffer():
     curbuf = vim.current.buffer
@@ -1586,14 +1589,15 @@ def submit_current_buffer():
     full_request = '\n'.join(curbuf)
 
     req = parse_request(full_request)
-    dest_host, dest_port, use_tls = dest_loc()
+    dest_host, dest_port, use_tls, storage_id = dest_loc()
     req.dest_host = dest_host
     req.dest_port = dest_port
     req.use_tls = use_tls
 
     comm_type, comm_addr = get_conn_addr()
     with ProxyConnection(kind=comm_type, addr=comm_addr) as conn:
-        new_req = conn.submit(req)
+        new_req = conn.submit(req, storage=storage_id)
+        conn.add_tag(new_req.db_id, "repeater", storage_id)
         update_buffers(new_req)
     
 # (left, right) = set_up_windows()

python/puppy/puppyproxy/interface/view.py

@@ -481,17 +481,23 @@ def site_map(client, args):
         paths = True
     else:
         paths = False
-    reqs = client.in_context_requests(headers_only=True)
-    paths_set = set()
-    for req in reqs:
-        if req.response and req.response.status_code != 404:
-            paths_set.add(path_tuple(req.url))
-    tree = sorted(list(paths_set))
-    if paths:
-        for p in tree:
-            print ('/'.join(list(p)))
-    else:
-        print_tree(tree)
+    all_reqs = client.in_context_requests(headers_only=True)
+    reqs_by_host = {}
+    for req in all_reqs:
+        reqs_by_host.setdefault(req.dest_host, []).append(req)
+    for host, reqs in reqs_by_host.items():
+        paths_set = set()
+        for req in reqs:
+            if req.response and req.response.status_code != 404:
+                paths_set.add(path_tuple(req.url))
+        tree = sorted(list(paths_set))
+        print(host)
+        if paths:
+            for p in tree:
+                print ('/'.join(list(p)))
+        else:
+            print_tree(tree)
+        print("")
 
 def dump_response(client, args):
     """
@@ -515,6 +521,78 @@ def dump_response(client, args):
     else:
         print('Request {} does not have a response'.format(req.reqid))
 
+def get_surrounding_lines(s, n, lines):
+    left = n
+    right = n
+    lines_left = 0
+    lines_right = 0
+
+    # move left until we find enough lines or hit the edge
+    while left > 0 and lines_left < lines:
+        if s[left] == '\n':
+            lines_left += 1
+        left -= 1
+
+    # move right until we find enough lines or hit the edge
+    while right < len(s) and lines_right < lines:
+        if s[right] == '\n':
+            lines_right += 1
+        right += 1
+
+    return s[left:right]
+
+def print_search_header(reqid, locstr):
+    printstr = Styles.TABLE_HEADER
+    printstr += "Result(s) for request {} ({})".format(reqid, locstr)
+    printstr += Colors.ENDC
+    print(printstr)
+
+def highlight_str(s, substr):
+    highlighted = Colors.BGYELLOW + Colors.BLACK + Colors.BOLD + substr + Colors.ENDC
+    return s.replace(substr, highlighted)
+
+def search_message(mes, substr, lines, reqid, locstr):
+    header_printed = False
+    for m in re.finditer(substr, mes):
+        if not header_printed:
+            print_search_header(reqid, locstr)
+            header_printed = True
+        n = m.start()
+        linestr = get_surrounding_lines(mes, n, lines)
+        linelist = linestr.split('\n')
+        linestr = '\n'.join(line[:500] for line in linelist)
+        toprint = highlight_str(linestr, substr)
+        print(toprint)
+        print('-'*50)
+
+def search(client, args):
+    search_str = args[0]
+    lines = 2
+    if len(args) > 1:
+        lines = int(args[1])
+    for req in client.in_context_requests_iter():
+        reqid = client.get_reqid(req)
+        reqheader_printed = False
+        try:
+            mes = req.full_message().decode()
+            search_message(mes, search_str, lines, reqid, "Request")
+        except UnicodeDecodeError:
+            pass
+        if req.response:
+            try:
+                mes = req.response.full_message().decode()
+                search_message(mes, search_str, lines, reqid, "Response")
+            except UnicodeDecodeError:
+                pass
+
+        wsheader_printed = False
+        for wsm in req.ws_messages:
+            if not wsheader_printed:
+                print_search_header(client.get_reqid(req), reqid, "Websocket Messages")
+                wsheader_printed = True
+            if search_str in wsm.message:
+                print(highlight_str(wsm.message, search_str))
+
 
 # @crochet.wait_for(timeout=None)
 # @defer.inlineCallbacks
@@ -572,6 +650,7 @@ def load_cmds(cmd):
         'urls': (find_urls, None),
         'site_map': (site_map, None),
         'dump_response': (dump_response, None),
+        'search': (search, None),
         # 'view_request_bytes': (view_request_bytes, None),
         # 'view_response_bytes': (view_response_bytes, None),
     })

python/puppy/puppyproxy/proxy.py

@@ -85,10 +85,16 @@ class SockBuffer:
 
 class Headers:
     def __init__(self, headers=None):
-        if headers is None:
-            self.headers = {}
-        else:
-            self.headers = headers
+        self.headers = {}
+        if headers is not None:
+            if isinstance(headers, Headers):
+                for _, pairs in headers.headers.items():
+                    for k, v in pairs:
+                        self.add(k, v)
+            else:
+                for k, vs in headers.items():
+                    for v in vs:
+                        self.add(k, v)
         
     def __contains__(self, hd):
         for k, _ in self.headers.items():
@@ -265,11 +271,7 @@ class HTTPRequest:
         self.proto_major = proto_major
         self.proto_minor = proto_minor
 
-        self.headers = Headers()
-        if headers is not None:
-            for k, vs in headers.items():
-                for v in vs:
-                    self.headers.add(k, v)
+        self.headers = Headers(headers)
         
         self.headers_only = headers_only
         self._body = bytes()
@@ -280,8 +282,8 @@ class HTTPRequest:
         self.dest_host = dest_host
         self.dest_port = dest_port
         self.use_tls = use_tls
-        self.time_start = time_start or datetime.datetime(1970, 1, 1)
-        self.time_end = time_end or datetime.datetime(1970, 1, 1)
+        self.time_start = time_start
+        self.time_end = time_end
         
         self.response = None
         self.unmangled = None
@@ -412,7 +414,7 @@ class HTTPRequest:
             path=self.url.geturl(),
             proto_major=self.proto_major,
             proto_minor=self.proto_minor,
-            headers=self.headers.headers,
+            headers=self.headers,
             body=self.body,
             dest_host=self.dest_host,
             dest_port=self.dest_port,
@@ -929,6 +931,21 @@ class ProxyConnection:
             ret.append(SavedStorage(ss["Id"], ss["Description"]))
         return ret
 
+    @messagingFunction
+    def set_proxy(self, use_proxy=False, proxy_host="", proxy_port=0, use_creds=False,
+            username="", password="", is_socks=False):
+        cmd = {
+            "Command": "SetProxy",
+            "UseProxy": use_proxy,
+            "ProxyHost": proxy_host,
+            "ProxyPort": proxy_port,
+            "ProxyIsSOCKS": is_socks,
+            "UseCredentials": use_creds,
+            "Username": username,
+            "Password": password,
+        }
+        self.reqrsp_cmd(cmd)
+        
     @messagingFunction
     def intercept(self, macro):
         # Run an intercepting macro until closed
@@ -1086,6 +1103,7 @@ class ProxyClient:
             # "add_in_memory_storage",
             # "close_storage",
             # "set_proxy_storage",
+            "set_proxy"
         }
         
     def __enter__(self):
@@ -1173,6 +1191,10 @@ class ProxyClient:
         storage = self.storage_by_prefix[prefix]
         return storage, realid
 
+    def get_reqid(self, req):
+        storage = self.storage_by_id[req.storage_id]
+        return storage.prefix + req.db_id
+        
     def storage_iter(self):
         for _, s in self.storage_by_id.items():
             yield s
@@ -1191,6 +1213,17 @@ class ProxyClient:
             ret = results[:max_results]
         return ret
 
+    def in_context_requests_iter(self, headers_only=False, max_results=0):
+        results = self.query_storage(self.context.query,
+                                     headers_only=headers_only,
+                                     max_results=max_results)
+        ret = results
+        if max_results > 0 and len(results) > max_results:
+            ret = results[:max_results]
+        for reqh in ret:
+            req = self.req_by_id(reqh.db_id, storage_id=reqh.storage_id)
+            yield req
+    
     def prefixed_reqid(self, req):
         prefix = ""
         if req.storage_id in self.storage_by_id:
@@ -1246,10 +1279,14 @@ class ProxyClient:
         results = [r for r in reversed(results)]
         return results
             
-    def req_by_id(self, reqid, headers_only=False):
-        storage, rid = self.parse_reqid(reqid)
-        return self.msg_conn.req_by_id(rid, headers_only=headers_only,
-                                       storage=storage.storage_id)
+    def req_by_id(self, reqid, storage_id=None, headers_only=False):
+        if storage_id is None:
+            storage, db_id = self.parse_reqid(reqid)
+            storage_id = storage.storage_id
+        else:
+            db_id = reqid
+        return self.msg_conn.req_by_id(db_id, headers_only=headers_only,
+                                       storage=storage_id)
 
     # for these and submit, might need storage stored on the request itself
     def add_tag(self, reqid, tag, storage=None):
@@ -1275,12 +1312,12 @@ class ProxyClient:
 
 
 def decode_req(result, headers_only=False):
-    if "StartTime" in result:
+    if "StartTime" in result and result["StartTime"] > 0:
         time_start = time_from_nsecs(result["StartTime"])
     else:
         time_start = None
 
-    if "EndTime" in result:
+    if "EndTime" in result and result["EndTime"] > 0:
         time_end = time_from_nsecs(result["EndTime"])
     else:
         time_end = None

python/puppy/puppyproxy/pup.py

@@ -114,6 +114,13 @@ def main():
                         client.add_listener(iface, port)
                     except MessageError as e:
                         print(str(e))
+
+                # Set upstream proxy
+                if config.use_proxy:
+                    client.set_proxy(config.use_proxy,
+                                     config.proxy_host,
+                                     config.proxy_port,
+                                     config.is_socks_proxy)
             interface_loop(client)
         except MessageError as e:
             print(str(e))

python/puppy/puppyproxy/util.py

@@ -2,6 +2,7 @@ import sys
 import string
 import time
 import datetime
+import base64
 from pygments.formatters import TerminalFormatter
 from pygments.lexers import get_lexer_for_mimetype, HttpLexer
 from pygments import highlight
@@ -275,8 +276,8 @@ def clipboard_contents():
 
 def encode_basic_auth(username, password):
     decoded = '%s:%s' % (username, password)
-    encoded = base64.b64encode(decoded)
-    header = 'Basic %s' % encoded
+    encoded = base64.b64encode(decoded.encode())
+    header = 'Basic %s' % encoded.decode()
     return header
 
 def parse_basic_auth(header):

schema.go

@@ -6,8 +6,8 @@ import (
 	"fmt"
 	"log"
 	"runtime"
-	"strings"
 	"sort"
+	"strings"
 )
 
 type schemaUpdater func(tx *sql.Tx) error
@@ -110,19 +110,19 @@ SCHEMA 8 / INITIAL
 func schema8(tx *sql.Tx) error {
 	// Create a schema that is the same as pappy's last version
 
-	cmds := []string {
+	cmds := []string{
 
-	`
+		`
     CREATE TABLE schema_meta (
             version INTEGER NOT NULL
         );
     `,
 
-	`
+		`
     INSERT INTO "schema_meta" VALUES(8);
     `,
 
-	`
+		`
     CREATE TABLE responses (
         id             INTEGER   PRIMARY KEY  AUTOINCREMENT,
         full_response  BLOB  NOT NULL,
@@ -130,28 +130,28 @@ func schema8(tx *sql.Tx) error {
     );
     `,
 
-	`
+		`
     CREATE TABLE scope (
             filter_order    INTEGER  NOT NULL,
             filter_string   TEXT  NOT NULL
         );
     `,
 
-	`
+		`
     CREATE TABLE tags (
             id       INTEGER   PRIMARY KEY  AUTOINCREMENT,
             tag      TEXT      NOT NULL
         );
     `,
 
-	`
+		`
     CREATE TABLE tagged (
             reqid       INTEGER,
             tagid       INTEGER
         );
     `,
 
-	`
+		`
     CREATE TABLE "requests" (
             id                INTEGER   PRIMARY KEY  AUTOINCREMENT,
             full_request      BLOB      NOT NULL,
@@ -167,7 +167,7 @@ func schema8(tx *sql.Tx) error {
         );
     `,
 
-	`
+		`
     CREATE TABLE saved_contexts (
             id                INTEGER   PRIMARY KEY  AUTOINCREMENT,
             context_name      TEXT      UNIQUE,
@@ -175,7 +175,7 @@ func schema8(tx *sql.Tx) error {
         );
     `,
 
-	`
+		`
     CREATE TABLE websocket_messages (
             id                INTEGER   PRIMARY KEY  AUTOINCREMENT,
             parent_request    INTEGER   REFERENCES requests(id),
@@ -187,7 +187,7 @@ func schema8(tx *sql.Tx) error {
         );
     `,
 
-	`
+		`
     CREATE INDEX ind_start_time ON requests(start_datetime);
     `,
 	}
@@ -240,7 +240,7 @@ func pappyListToStrMessageQuery(f []string) (StrMessageQuery, error) {
 }
 
 type s9ScopeStr struct {
-	Order int64
+	Order  int64
 	Filter string
 }
 
@@ -260,8 +260,8 @@ func (ls s9ScopeSort) Less(i int, j int) bool {
 
 func schema9(tx *sql.Tx) error {
 	/*
-     Converts the floating point timestamps into integers representing nanoseconds from jan 1 1970
-     */
+	   Converts the floating point timestamps into integers representing nanoseconds from jan 1 1970
+	*/
 
 	// Rename the old requests table
 	if err := execute(tx, "ALTER TABLE requests RENAME TO requests_old"); err != nil {
@@ -289,13 +289,13 @@ func schema9(tx *sql.Tx) error {
         );
     `,
 
-    `
+		`
     INSERT INTO requests
     SELECT id, full_request, submitted, response_id, unmangled_id, port, is_ssl, host, plugin_data, 0, 0
     FROM requests_old
     `,
 
-	`
+		`
     CREATE TABLE websocket_messages (
             id                INTEGER   PRIMARY KEY  AUTOINCREMENT,
             parent_request    INTEGER   REFERENCES requests(id),
@@ -307,7 +307,7 @@ func schema9(tx *sql.Tx) error {
         );
     `,
 
-	`
+		`
     INSERT INTO websocket_messages
     SELECT id, parent_request, unmangled_id, is_binary, direction, 0, contents
     FROM websocket_messages_old
@@ -337,13 +337,13 @@ func schema9(tx *sql.Tx) error {
 
 		if startDT.Valid {
 			// Convert to nanoseconds
-			newStartDT = int64(startDT.Float64*1000000000)
+			newStartDT = int64(startDT.Float64 * 1000000000)
 		} else {
 			newStartDT = 0
 		}
 
 		if endDT.Valid {
-			newEndDT = int64(endDT.Float64*1000000000)
+			newEndDT = int64(endDT.Float64 * 1000000000)
 		} else {
 			newEndDT = 0
 		}
@@ -378,7 +378,7 @@ func schema9(tx *sql.Tx) error {
 
 		if sentDT.Valid {
 			// Convert to nanoseconds
-			newSentDT = int64(startDT.Float64*1000000000)
+			newSentDT = int64(startDT.Float64 * 1000000000)
 		} else {
 			newSentDT = 0
 		}

search.go

@@ -68,12 +68,12 @@ const (
 
 // A struct representing the data to be searched for a pair such as a header or url param
 type PairValue struct {
-	key string
+	key   string
 	value string
 }
 
 type QueryPhrase [][]interface{} // A list of queries. Will match if any queries match the request
-type MessageQuery []QueryPhrase // A list of phrases. Will match if all the phrases match the request
+type MessageQuery []QueryPhrase  // A list of phrases. Will match if all the phrases match the request
 
 type StrQueryPhrase [][]string
 type StrMessageQuery []StrQueryPhrase
@@ -451,7 +451,7 @@ func pairValuesFromCookies(cookies []*http.Cookie) []*PairValue {
 	return pairs
 }
 
-func pairsToStrings(pairs []*PairValue) ([]string) {
+func pairsToStrings(pairs []*PairValue) []string {
 	// Converts a list of pairs into a list of strings containing all keys and values
 	// k1: v1, k2: v2 -> ["k1", "v1", "k2", "v2"]
 	strs := make([]string, 0)
@@ -710,9 +710,9 @@ func FieldStrToGo(field string) (SearchField, error) {
 		return FieldBothCookie, nil
 	case "tag":
 		return FieldTag, nil
-	case "after":
+	case "after", "af":
 		return FieldAfter, nil
-	case "before":
+	case "before", "b4":
 		return FieldBefore, nil
 	case "timerange":
 		return FieldTimeRange, nil

search_test.go

@@ -8,7 +8,9 @@ import (
 
 func checkSearch(t *testing.T, req *ProxyRequest, expected bool, args ...interface{}) {
 	checker, err := NewRequestChecker(args...)
-	if err != nil { t.Error(err.Error()) }
+	if err != nil {
+		t.Error(err.Error())
+	}
 	result := checker(req)
 	if result != expected {
 		_, f, ln, _ := runtime.Caller(1)
@@ -18,9 +20,13 @@ func checkSearch(t *testing.T, req *ProxyRequest, expected bool, args ...interfa
 
 func TestAllSearch(t *testing.T) {
 	checker, err := NewRequestChecker(FieldAll, StrContains, "foo")
-	if err != nil { t.Error(err.Error()) }
+	if err != nil {
+		t.Error(err.Error())
+	}
 	req := testReq()
-	if !checker(req) { t.Error("Failed to match FieldAll, StrContains") }
+	if !checker(req) {
+		t.Error("Failed to match FieldAll, StrContains")
+	}
 }
 
 func TestBodySearch(t *testing.T) {

signer.go

@@ -190,4 +190,3 @@ func SignHost(ca tls.Certificate, hosts []string) (cert tls.Certificate, err err
 		PrivateKey:  certpriv,
 	}, nil
 }
-

sqlitestorage.go

@@ -11,17 +11,19 @@ import (
 	"sync"
 	"time"
 
-	_ "github.com/mattn/go-sqlite3"
 	"github.com/gorilla/websocket"
+	_ "github.com/mattn/go-sqlite3"
 )
 
 var REQUEST_SELECT string = "SELECT id, full_request, response_id, unmangled_id, port, is_ssl, host, start_datetime, end_datetime FROM requests"
 var RESPONSE_SELECT string = "SELECT id, full_response, unmangled_id FROM responses"
 var WS_SELECT string = "SELECT id, parent_request, unmangled_id, is_binary, direction, time_sent, contents FROM websocket_messages"
 
+var inmemIdCounter = IdCounter()
+
 type SQLiteStorage struct {
 	dbConn *sql.DB
-	mtx sync.Mutex
+	mtx    sync.Mutex
 	logger *log.Logger
 }
 
@@ -49,7 +51,8 @@ func OpenSQLiteStorage(fname string, logger *log.Logger) (*SQLiteStorage, error)
 }
 
 func InMemoryStorage(logger *log.Logger) (*SQLiteStorage, error) {
-	return OpenSQLiteStorage("file::memory:?mode=memory&cache=shared", logger)
+	var toOpen = fmt.Sprintf("file:inmem%d:memory:?mode=memory&cache=shared", inmemIdCounter())
+	return OpenSQLiteStorage(toOpen, logger)
 }
 
 func (rs *SQLiteStorage) Close() {
@@ -272,7 +275,7 @@ func wsFromRow(tx *sql.Tx, ms *SQLiteStorage, id sql.NullInt64, parent_request s
 	return wsm, nil
 }
 
-func addTagsToStorage(tx *sql.Tx, req *ProxyRequest) (error) {
+func addTagsToStorage(tx *sql.Tx, req *ProxyRequest) error {
 	// Save the tags
 	for _, tag := range req.Tags() {
 		var db_tagid sql.NullInt64
@@ -378,7 +381,6 @@ func (ms *SQLiteStorage) saveNewRequest(tx *sql.Tx, req *ProxyRequest) error {
 	var rspid *string
 	var unmangledId *string
 
-
 	if req.ServerResponse != nil {
 		if req.ServerResponse.DbId == "" {
 			return errors.New("response has not been saved yet, cannot save request")
@@ -557,20 +559,20 @@ func (ms *SQLiteStorage) loadRequest(tx *sql.Tx, reqid string) (*ProxyRequest, e
 	var db_end_datetime sql.NullInt64
 
 	// err = tx.QueryRow(`
-    //     SELECT
-    //     id, full_request, response_id, unmangled_id, port, is_ssl, host, start_datetime, end_datetime
-    //     FROM requests WHERE id=?`, dbId).Scan(
-	err = tx.QueryRow(REQUEST_SELECT + " WHERE id=?", dbId).Scan(
-			&db_id,
-			&db_full_request,
-			&db_response_id,
-			&db_unmangled_id,
-			&db_port,
-			&db_is_ssl,
-			&db_host,
-			&db_start_datetime,
-			&db_end_datetime,
-		)
+	//     SELECT
+	//     id, full_request, response_id, unmangled_id, port, is_ssl, host, start_datetime, end_datetime
+	//     FROM requests WHERE id=?`, dbId).Scan(
+	err = tx.QueryRow(REQUEST_SELECT+" WHERE id=?", dbId).Scan(
+		&db_id,
+		&db_full_request,
+		&db_response_id,
+		&db_unmangled_id,
+		&db_port,
+		&db_is_ssl,
+		&db_host,
+		&db_start_datetime,
+		&db_end_datetime,
+	)
 	if err == sql.ErrNoRows {
 		return nil, fmt.Errorf("Request with id %d does not exist", dbId)
 	} else if err != nil {
@@ -624,7 +626,7 @@ func (ms *SQLiteStorage) loadUnmangledRequest(tx *sql.Tx, reqid string) (*ProxyR
 	return ms.loadRequest(tx, strconv.FormatInt(db_unmangled_id.Int64, 10))
 }
 
-func (ms *SQLiteStorage) DeleteRequest(reqid string) (error) {
+func (ms *SQLiteStorage) DeleteRequest(reqid string) error {
 	ms.mtx.Lock()
 	defer ms.mtx.Unlock()
 	tx, err := ms.dbConn.Begin()
@@ -640,7 +642,7 @@ func (ms *SQLiteStorage) DeleteRequest(reqid string) (error) {
 	return nil
 }
 
-func (ms *SQLiteStorage) deleteRequest(tx *sql.Tx, reqid string) (error) {
+func (ms *SQLiteStorage) deleteRequest(tx *sql.Tx, reqid string) error {
 	if reqid == "" {
 		return nil
 	}
@@ -842,11 +844,11 @@ func (ms *SQLiteStorage) loadResponse(tx *sql.Tx, rspid string) (*ProxyResponse,
 	var db_full_response []byte
 	var db_unmangled_id sql.NullInt64
 
-	err = tx.QueryRow(RESPONSE_SELECT + " WHERE id=?", dbId).Scan(
-			&db_id,
-			&db_full_response,
-			&db_unmangled_id,
-		)
+	err = tx.QueryRow(RESPONSE_SELECT+" WHERE id=?", dbId).Scan(
+		&db_id,
+		&db_full_response,
+		&db_unmangled_id,
+	)
 	if err == sql.ErrNoRows {
 		return nil, fmt.Errorf("Response with id %d does not exist", dbId)
 	} else if err != nil {
@@ -1141,16 +1143,15 @@ func (ms *SQLiteStorage) loadWSMessage(tx *sql.Tx, wsmid string) (*ProxyWSMessag
 	var db_time_sent sql.NullInt64
 	var db_contents []byte
 
-
-	err = tx.QueryRow(WS_SELECT + " WHERE id=?", dbId).Scan(
-			&db_id,
-			&db_parent_request,
-			&db_unmangled_id,
-			&db_is_binary,
-			&db_direction,
-			&db_time_sent,
-			&db_contents,
-		)
+	err = tx.QueryRow(WS_SELECT+" WHERE id=?", dbId).Scan(
+		&db_id,
+		&db_parent_request,
+		&db_unmangled_id,
+		&db_is_binary,
+		&db_direction,
+		&db_time_sent,
+		&db_contents,
+	)
 	if err == sql.ErrNoRows {
 		return nil, fmt.Errorf("Message with id %d does not exist", dbId)
 	} else if err != nil {
@@ -1432,7 +1433,7 @@ func (ms *SQLiteStorage) checkRequests(tx *sql.Tx, limit int64, checker RequestC
 	return ms.reqSearchHelper(tx, limit, checker, "")
 }
 
-func (ms *SQLiteStorage) SaveQuery(name string, query MessageQuery) (error) {
+func (ms *SQLiteStorage) SaveQuery(name string, query MessageQuery) error {
 	ms.mtx.Lock()
 	defer ms.mtx.Unlock()
 	tx, err := ms.dbConn.Begin()
@@ -1448,7 +1449,7 @@ func (ms *SQLiteStorage) SaveQuery(name string, query MessageQuery) (error) {
 	return nil
 }
 
-func (ms *SQLiteStorage) saveQuery(tx *sql.Tx, name string, query MessageQuery) (error) {
+func (ms *SQLiteStorage) saveQuery(tx *sql.Tx, name string, query MessageQuery) error {
 	strQuery, err := GoQueryToStrQuery(query)
 	if err != nil {
 		return fmt.Errorf("error creating string version of query: %s", err.Error())
@@ -1523,7 +1524,7 @@ func (ms *SQLiteStorage) loadQuery(tx *sql.Tx, name string) (MessageQuery, error
 	return retQuery, nil
 }
 
-func (ms *SQLiteStorage) DeleteQuery(name string) (error) {
+func (ms *SQLiteStorage) DeleteQuery(name string) error {
 	ms.mtx.Lock()
 	defer ms.mtx.Unlock()
 	tx, err := ms.dbConn.Begin()
@@ -1539,7 +1540,7 @@ func (ms *SQLiteStorage) DeleteQuery(name string) (error) {
 	return nil
 }
 
-func (ms *SQLiteStorage) deleteQuery(tx *sql.Tx, name string) (error) {
+func (ms *SQLiteStorage) deleteQuery(tx *sql.Tx, name string) error {
 	stmt, err := tx.Prepare(`DELETE FROM saved_contexts WHERE context_name=?;`)
 	if err != nil {
 		return fmt.Errorf("error preparing statement to insert request into database: %s", err.Error())
@@ -1602,4 +1603,3 @@ func (ms *SQLiteStorage) allSavedQueries(tx *sql.Tx) ([]*SavedQuery, error) {
 	}
 	return savedQueries, nil
 }
-

sqlitestorage_test.go

@@ -1,10 +1,9 @@
 package main
 
 import (
-	"testing"
 	"runtime"
+	"testing"
 	"time"
-	"fmt"
 )
 
 func testStorage() *SQLiteStorage {
@@ -20,7 +19,7 @@ func checkTags(t *testing.T, result, expected []string) {
 		return
 	}
 
-	for i, a := range(result) {
+	for i, a := range result {
 		b := expected[i]
 		if a != b {
 			t.Errorf("Failed tag test at %s:%d. Expected %s, got %s", f, ln, expected, result)
@@ -56,7 +55,6 @@ func TestTagging(t *testing.T) {
 	checkTags(t, req3.Tags(), []string{"bar"})
 }
 
-
 func TestTime(t *testing.T) {
 	req := testReq()
 	req.StartDatetime = time.Unix(0, 1234567)

storage.go

@@ -1,8 +1,8 @@
 package main
 
 import (
-	"fmt"
 	"errors"
+	"fmt"
 )
 
 type MessageStorage interface {
@@ -20,7 +20,7 @@ type MessageStorage interface {
 	LoadRequest(reqid string) (*ProxyRequest, error)
 	LoadUnmangledRequest(reqid string) (*ProxyRequest, error)
 	// Delete a request
-	DeleteRequest(reqid string) (error)
+	DeleteRequest(reqid string) error
 
 	// Update an existing response in the storage. Requires that it has already been saved
 	UpdateResponse(rsp *ProxyResponse) error
@@ -30,7 +30,7 @@ type MessageStorage interface {
 	LoadResponse(rspid string) (*ProxyResponse, error)
 	LoadUnmangledResponse(rspid string) (*ProxyResponse, error)
 	// Delete a response
-	DeleteResponse(rspid string) (error)
+	DeleteResponse(rspid string) error
 
 	// Update an existing websocket message in the storage. Requires that it has already been saved
 	UpdateWSMessage(req *ProxyRequest, wsm *ProxyWSMessage) error
@@ -40,7 +40,7 @@ type MessageStorage interface {
 	LoadWSMessage(wsmid string) (*ProxyWSMessage, error)
 	LoadUnmangledWSMessage(wsmid string) (*ProxyWSMessage, error)
 	// Delete a WSMessage
-	DeleteWSMessage(wsmid string) (error)
+	DeleteWSMessage(wsmid string) error
 
 	// Get list of all the request keys
 	RequestKeys() ([]string, error)
@@ -57,9 +57,9 @@ type MessageStorage interface {
 
 	// Query functions
 	AllSavedQueries() ([]*SavedQuery, error)
-	SaveQuery(name string, query MessageQuery) (error)
+	SaveQuery(name string, query MessageQuery) error
 	LoadQuery(name string) (MessageQuery, error)
-	DeleteQuery(name string) (error)
+	DeleteQuery(name string) error
 }
 
 const QueryNotSupported = ConstErr("custom query not supported")
@@ -67,7 +67,7 @@ const QueryNotSupported = ConstErr("custom query not supported")
 type ReqSort []*ProxyRequest
 
 type SavedQuery struct {
-	Name string
+	Name  string
 	Query MessageQuery
 }
 
@@ -118,7 +118,7 @@ func SaveNewRequest(ms MessageStorage, req *ProxyRequest) error {
 	}
 
 	if err := ms.SaveNewRequest(req); err != nil {
-			return fmt.Errorf("error saving new request: %s", err.Error())
+		return fmt.Errorf("error saving new request: %s", err.Error())
 	}
 
 	for _, wsm := range req.WSMessages {
@@ -224,4 +224,3 @@ func UpdateWSMessage(ms MessageStorage, req *ProxyRequest, wsm *ProxyWSMessage)
 		return ms.UpdateWSMessage(req, wsm)
 	}
 }
-

testutil.go

@@ -1,11 +1,11 @@
 package main
 
 import (
-	"testing"
 	"runtime"
+	"testing"
 )
 
-func testReq() (*ProxyRequest) {
+func testReq() *ProxyRequest {
 	testReq, _ := ProxyRequestFromBytes(
 		[]byte("POST /?foo=bar HTTP/1.1\r\nFoo: Bar\r\nCookie: cookie=choco\r\nContent-Length: 7\r\n\r\nfoo=baz"),
 		"foobaz",

util.go

@@ -1,16 +1,16 @@
 package main
 
 import (
-	"sync"
-	"log"
 	"io/ioutil"
+	"log"
+	"sync"
 )
 
 type ConstErr string
 
 func (e ConstErr) Error() string { return string(e) }
 
-func DuplicateBytes(bs []byte) ([]byte) {
+func DuplicateBytes(bs []byte) []byte {
 	retBs := make([]byte, len(bs))
 	copy(retBs, bs)
 	return retBs

webui.go

@@ -0,0 +1,169 @@
+package main
+
+import (
+	"encoding/pem"
+	"html/template"
+	"net/http"
+	"strings"
+)
+
+// Page template
+var MASTER_SRC string = `
+<html>
+<head>
+<title>{{block "title" .}}Puppy Proxy{{end}}</title>
+{{block "head" .}}{{end}}
+</head>
+<body>
+{{block "body" .}}{{end}}
+</body>
+</html>
+`
+var MASTER_TPL *template.Template
+
+// Page sources
+var HOME_SRC string = `
+{{define "title"}}Puppy Home{{end}}
+{{define "body"}}
+	<p>Welcome to Puppy<p>
+	<ul>
+	<li><a href="/certs">Download CA certificate</a></li>
+	</ul>
+{{end}}
+`
+var HOME_TPL *template.Template
+
+var CERTS_SRC string = `
+{{define "title"}}CA Certificate{{end}}
+{{define "body"}}
+	<p>Downlad this CA cert and add it to your browser to intercept HTTPS messages<p>
+	<p><a href="/certs/download">Download</p>
+{{end}}
+`
+var CERTS_TPL *template.Template
+
+var RSPVIEW_SRC string = `
+{{define "title"}}Response Viewer{{end}}
+{{define "head"}}
+	<script>
+	function ViewResponse() {
+		rspid = document.getElementById("rspid").value
+		window.location.href = "/rsp/" + rspid
+	}
+	</script>
+{{end}}
+{{define "body"}}
+	<p>Enter a response ID below to view it in the browser<p>
+	<input type="text" id="rspid"></input><input type="button" onclick="ViewResponse()" value="Go!"></input>
+{{end}}
+`
+var RSPVIEW_TPL *template.Template
+
+func init() {
+	var err error
+	MASTER_TPL, err = template.New("master").Parse(MASTER_SRC)
+	if err != nil {
+		panic(err)
+	}
+
+	HOME_TPL, err = template.Must(MASTER_TPL.Clone()).Parse(HOME_SRC)
+	if err != nil {
+		panic(err)
+	}
+
+	CERTS_TPL, err = template.Must(MASTER_TPL.Clone()).Parse(CERTS_SRC)
+	if err != nil {
+		panic(err)
+	}
+
+	RSPVIEW_TPL, err = template.Must(MASTER_TPL.Clone()).Parse(RSPVIEW_SRC)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func responseHeaders(w http.ResponseWriter) {
+	w.Header().Set("Connection", "close")
+	w.Header().Set("Cache-control", "no-cache")
+	w.Header().Set("Pragma", "no-cache")
+	w.Header().Set("Cache-control", "no-store")
+	w.Header().Set("X-Frame-Options", "DENY")
+}
+
+func WebUIHandler(w http.ResponseWriter, r *http.Request, iproxy *InterceptingProxy) {
+	responseHeaders(w)
+	parts := strings.Split(r.URL.Path, "/")
+	switch parts[1] {
+	case "":
+		WebUIRootHandler(w, r, iproxy)
+	case "certs":
+		WebUICertsHandler(w, r, iproxy, parts[2:])
+	case "rsp":
+		WebUIRspHandler(w, r, iproxy, parts[2:])
+	}
+}
+
+func WebUIRootHandler(w http.ResponseWriter, r *http.Request, iproxy *InterceptingProxy) {
+	err := HOME_TPL.Execute(w, nil)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+}
+
+func WebUICertsHandler(w http.ResponseWriter, r *http.Request, iproxy *InterceptingProxy, path []string) {
+	if len(path) > 0 && path[0] == "download" {
+		cert := iproxy.GetCACertificate()
+		if cert == nil {
+			w.Write([]byte("no active certs to download"))
+			return
+		}
+
+		pemData := pem.EncodeToMemory(
+			&pem.Block{
+				Type:  "CERTIFICATE",
+				Bytes: cert.Certificate[0],
+			},
+		)
+		w.Header().Set("Content-Type", "application/octet-stream")
+		w.Header().Set("Content-Disposition", "attachment; filename=\"cert.pem\"")
+		w.Write(pemData)
+		return
+	}
+	err := CERTS_TPL.Execute(w, nil)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+}
+
+func viewResponseHeaders(w http.ResponseWriter) {
+	w.Header().Del("Cookie")
+}
+
+func WebUIRspHandler(w http.ResponseWriter, r *http.Request, iproxy *InterceptingProxy, path []string) {
+	if len(path) > 0 {
+		reqid := path[0]
+		ms := iproxy.GetProxyStorage()
+		req, err := ms.LoadRequest(reqid)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		rsp := req.ServerResponse
+		for k, v := range rsp.Header {
+			for _, vv := range v {
+				w.Header().Add(k, vv)
+			}
+		}
+		viewResponseHeaders(w)
+		w.WriteHeader(rsp.StatusCode)
+		w.Write(rsp.BodyBytes())
+		return
+	}
+	err := RSPVIEW_TPL.Execute(w, nil)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+}