From aea6c6d2f03db4aff73ec78cf0fd34ae6b327df1 Mon Sep 17 00:00:00 2001
From: Anders Pitman <tapitman11@gmail.com>
Date: Thu, 10 Dec 2020 23:14:34 -0700
Subject: [PATCH] Make client easier to import into other code

---
 client.go               | 43 ++++++++++++++++++++---------------------
 cmd/boringproxy/main.go | 24 ++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/client.go b/client.go
index 7d675e9..1858bd1 100644
--- a/client.go
+++ b/client.go
@@ -5,7 +5,6 @@ import (
 	"crypto/tls"
 	"encoding/json"
 	"errors"
-	"flag"
 	"fmt"
 	"github.com/caddyserver/certmagic"
 	"golang.org/x/crypto/ssh"
@@ -14,7 +13,6 @@ import (
 	"log"
 	"net"
 	"net/http"
-	"os"
 	"strings"
 	"sync"
 	"time"
@@ -33,25 +31,26 @@ type Client struct {
 	certConfig       *certmagic.Config
 }
 
-func NewClient() *Client {
-	flagSet := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
-	server := flagSet.String("server", "", "boringproxy server")
-	token := flagSet.String("token", "", "Access token")
-	name := flagSet.String("client-name", "", "Client name")
-	user := flagSet.String("user", "admin", "user")
-	certDir := flagSet.String("cert-dir", "", "TLS cert directory")
-	acmeEmail := flagSet.String("acme-email", "", "Email for ACME (ie Let's Encrypt)")
-	dnsServer := flagSet.String("dns-server", "", "Custom DNS server")
-	flagSet.Parse(os.Args[2:])
+type ClientConfig struct {
+	ServerAddr string
+	Token      string
+	ClientName string
+	User       string
+	CertDir    string
+	AcmeEmail  string
+	DnsServer  string
+}
 
-	if *dnsServer != "" {
+func NewClient(config *ClientConfig) *Client {
+
+	if config.DnsServer != "" {
 		net.DefaultResolver = &net.Resolver{
 			PreferGo: true,
 			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
 				d := net.Dialer{
 					Timeout: time.Millisecond * time.Duration(10000),
 				}
-				return d.DialContext(ctx, "udp", fmt.Sprintf("%s:53", *dnsServer))
+				return d.DialContext(ctx, "udp", fmt.Sprintf("%s:53", config.DnsServer))
 			},
 		}
 	}
@@ -71,12 +70,12 @@ func NewClient() *Client {
 
 	certmagic.DefaultACME.DisableHTTPChallenge = true
 
-	if *certDir != "" {
-		certmagic.Default.Storage = &certmagic.FileStorage{*certDir}
+	if config.CertDir != "" {
+		certmagic.Default.Storage = &certmagic.FileStorage{config.CertDir}
 	}
 
-	if *acmeEmail != "" {
-		certmagic.DefaultACME.Email = *acmeEmail
+	if config.AcmeEmail != "" {
+		certmagic.DefaultACME.Email = config.AcmeEmail
 	}
 
 	certConfig := certmagic.NewDefault()
@@ -90,10 +89,10 @@ func NewClient() *Client {
 		httpClient:       httpClient,
 		tunnels:          tunnels,
 		previousEtag:     "",
-		server:           *server,
-		token:            *token,
-		clientName:       *name,
-		user:             *user,
+		server:           config.ServerAddr,
+		token:            config.Token,
+		clientName:       config.ClientName,
+		user:             config.User,
 		cancelFuncs:      cancelFuncs,
 		cancelFuncsMutex: cancelFuncsMutex,
 		certConfig:       certConfig,
diff --git a/cmd/boringproxy/main.go b/cmd/boringproxy/main.go
index f0788ed..7f21c6d 100644
--- a/cmd/boringproxy/main.go
+++ b/cmd/boringproxy/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"github.com/boringproxy/boringproxy"
 	"log"
@@ -22,7 +23,28 @@ func main() {
 		boringproxy.Listen()
 
 	case "client":
-		client := boringproxy.NewClient()
+
+		flagSet := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
+		server := flagSet.String("server", "", "boringproxy server")
+		token := flagSet.String("token", "", "Access token")
+		name := flagSet.String("client-name", "", "Client name")
+		user := flagSet.String("user", "admin", "user")
+		certDir := flagSet.String("cert-dir", "", "TLS cert directory")
+		acmeEmail := flagSet.String("acme-email", "", "Email for ACME (ie Let's Encrypt)")
+		dnsServer := flagSet.String("dns-server", "", "Custom DNS server")
+		flagSet.Parse(os.Args[2:])
+
+		config := &boringproxy.ClientConfig{
+			ServerAddr: *server,
+			Token:      *token,
+			ClientName: *name,
+			User:       *user,
+			CertDir:    *certDir,
+			AcmeEmail:  *acmeEmail,
+			DnsServer:  *dnsServer,
+		}
+
+		client := boringproxy.NewClient(config)
 		client.RunPuppetClient()
 	default:
 		fmt.Println("Invalid command " + command)