1 files changed, 84 insertions, 0 deletions

diff --git a/common/pt_extras/pt_extras.go b/common/pt_extras/pt_extras.go
index 63c1aba..e38d44e 100644
--- a/common/pt_extras/pt_extras.go
+++ b/common/pt_extras/pt_extras.go
@@ -30,6 +30,10 @@ package pt_extras
 import (
 	"errors"
 	"fmt"
+	"github.com/OperatorFoundation/shapeshifter-dispatcher/common/log"
+	"github.com/OperatorFoundation/shapeshifter-transports/transports/Optimizer"
+	"github.com/OperatorFoundation/shapeshifter-transports/transports/obfs4"
+	"github.com/OperatorFoundation/shapeshifter-transports/transports/shadow"
 	"net"
 	"net/url"
 	"os"
@@ -170,3 +174,83 @@ func resolveAddrStr(addrStr string) (*net.TCPAddr, error) {
 func PtShouldExitOnStdinClose() bool {
 	return os.Getenv("TOR_PT_EXIT_ON_STDIN_CLOSE") == "1"
 }
+
+func ArgsToDialer(target string, name string, args pt.Args) (Optimizer.Transport, error) {
+	switch name {
+	//case "obfs2":
+	//	transport := obfs2.NewObfs2Transport()
+	//	dialer = transport.Dial
+	//	return dialer, nil
+	case "obfs4":
+		if cert, ok := args["cert"]; ok {
+			if iatModeStr, ok2 := args["iatMode"]; ok2 {
+				iatMode, err := strconv.Atoi(iatModeStr[0])
+				if err == nil {
+					transport := obfs4.Transport{
+						CertString: cert[0],
+						IatMode:    iatMode,
+						Address:    target,
+					}
+					return transport, nil
+				} else {
+					log.Errorf("obfs4 transport bad iatMode value: %s %s", iatModeStr[0], err)
+					return nil, errors.New("obfs4 transport bad iatMode value")
+				}
+			} else {
+				log.Errorf("obfs4 transport missing iatMode argument: %s", args)
+				return nil, errors.New("obfs4 transport missing iatMode argument")
+			}
+		} else {
+			log.Errorf("obfs4 transport missing cert argument: %s", args)
+			return nil, errors.New("obfs4 transport missing cert argument")
+		}
+	case "shadow":
+		if password, ok := args["password"]; ok {
+			if cipher, ok2 := args["cipherName"]; ok2 {
+				transport := shadow.Transport{
+					Password:   password[0],
+					CipherName: cipher[0],
+					Address:    target,
+				}
+				return transport, nil
+			} else {
+				log.Errorf("shadow transport missing cipher argument: %s", args)
+				return nil, errors.New("shadow transport missing cipher argument")
+			}
+		} else {
+			log.Errorf("shadow transport missing password argument: %s", args)
+			return nil, errors.New("shadow transport missing password argument")
+		}
+	case "Optimizer":
+		if _, ok := args["transports"]; ok {
+			if strategyName, ok2 := args["strategy"]; ok2 {
+				var strategy Optimizer.Strategy = nil
+				switch strategyName[0] {
+				case "first":
+					strategy = Optimizer.NewFirstStrategy()
+				case "random":
+					strategy = Optimizer.NewRandomStrategy()
+				case "rotate":
+					strategy = Optimizer.NewRotateStrategy()
+				case "track":
+					strategy = Optimizer.NewTrackStrategy()
+				case "min":
+					strategy = Optimizer.NewMinimizeDialDuration()
+				}
+				transports := []Optimizer.Transport{}
+				transport := Optimizer.NewOptimizerClient(transports, strategy)
+				return transport, nil
+			} else {
+				log.Errorf("Optimizer transport missing transports argument: %s", args)
+				return nil, errors.New("optimizer transport missing transports argument")
+			}
+		} else {
+			log.Errorf("Optimizer transport missing strategy argument: %s", args)
+			return nil, errors.New("optimizer transport missing strategy argument")
+		}
+
+	default:
+		log.Errorf("Unknown transport: %s", name)
+		return nil, errors.New("unknown transport")
+	}
+}
\ No newline at end of file