From 1d819cf78f065d20b4756bf9a7707e03d423f015 Mon Sep 17 00:00:00 2001
From: Jille Timmermans <jille@snoozethis.com>
Date: Mon, 23 Jan 2023 16:12:36 +0100
Subject: [PATCH] cli: Add cleaner flags for specifying where to connect to

---
 README.md          |  6 ++----
 cmd/pgperms/cli.go | 33 +++++++++++++++++++++++++++++----
 go.mod             |  3 +++
 go.sum             |  6 ++++++
 4 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 2af1877..a76464a 100644
--- a/README.md
+++ b/README.md
@@ -15,15 +15,13 @@ or grab it from our [Releases](https://github.com/SnoozeThis-org/pgperms/release
 If you already have an existing PostgreSQL cluster running, you can create pgperms config file from that database through:
 
 ```shell
-$ export DSN="host=localhost username=postgres dbname=postgres"
-$ pgperms --dump > myconfig.yaml
+$ pgperms --dump --user postgres --database=postgres > pgperms.yaml
 ```
 
 Then you can edit your config and bring PostgreSQL to the new desired state:
 
 ```shell
-$ export DSN="host=localhost username=postgres dbname=postgres"
-$ pgperms myconfig.yaml
+$ pgperms --user postgres --database=postgres --config pgperms.yaml
 ```
 
 ## Managing roles
diff --git a/cmd/pgperms/cli.go b/cmd/pgperms/cli.go
index 86570b1..7b72d9f 100644
--- a/cmd/pgperms/cli.go
+++ b/cmd/pgperms/cli.go
@@ -6,15 +6,25 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"strings"
 
 	"github.com/SnoozeThis-org/pgperms"
+	"github.com/creachadair/getpass"
 	"github.com/jackc/pgx/v4"
 	"github.com/spf13/pflag"
 )
 
 var (
-	dump = pflag.Bool("dump", false, "Whether to dump the current permissions")
+	defaultConfig, _ = pgx.ParseConfig("")
+
+	config      = pflag.StringP("config", "c", "pgperms.yaml", "Path to the pgperms yaml config file")
+	dump        = pflag.Bool("dump", false, "Whether to dump the current permissions")
 	showVersion = pflag.Bool("version", false, "Dump the version and exit")
+	host        = pflag.StringP("host", "h", defaultConfig.Host, "database server host or socket directory")
+	port        = pflag.IntP("port", "P", int(defaultConfig.Port), "database server port")
+	username    = pflag.StringP("username", "U", defaultConfig.User, "database user name")
+	askPassword = pflag.BoolP("password", "W", false, "prompt for password")
+	database    = pflag.StringP("database", "d", "postgres", "database name for initial connection")
 
 	// Injected by releaser
 	version string
@@ -31,7 +41,15 @@ func main() {
 		return
 	}
 	ctx := context.Background()
-	conn, err := pgx.Connect(ctx, os.Getenv("DSN"))
+	dsn := fmt.Sprintf("host=%s port=%d user=%s dbname=%s", escapeDSNString(*host), *port, escapeDSNString(*username), escapeDSNString(*database))
+	if *askPassword {
+		pass, err := getpass.Prompt("Password: ")
+		if err != nil {
+			log.Fatalf("Failed to read password from prompt: %v", err)
+		}
+		dsn += " password=" + escapeDSNString(pass)
+	}
+	conn, err := pgx.Connect(ctx, dsn)
 	if err != nil {
 		log.Fatalf("Failed to connect to database: %v", err)
 	}
@@ -43,9 +61,12 @@ func main() {
 		fmt.Println(ret)
 		return
 	}
-	desired, err := ioutil.ReadFile(pflag.Arg(0))
+	if *config == "" {
+		log.Fatalf("Unless --dump is specified, --config must be set")
+	}
+	desired, err := ioutil.ReadFile(*config)
 	if err != nil {
-		log.Fatalf("Failed to read from config file %q: %v", pflag.Arg(0), err)
+		log.Fatalf("Failed to read from config file %q: %v", *config, err)
 	}
 	rec := pgperms.NewRecorder()
 	if err := pgperms.Sync(ctx, pgperms.NewConnections(ctx, conn), desired, rec); err != nil {
@@ -55,3 +76,7 @@ func main() {
 		fmt.Println(q)
 	}
 }
+
+func escapeDSNString(s string) string {
+	return "'" + strings.ReplaceAll(strings.ReplaceAll(s, `\`, `\\`), `'`, `\'`) + "'"
+}
diff --git a/go.mod b/go.mod
index 64ab0c4..6dfcdc2 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,7 @@ go 1.18
 
 require (
 	github.com/Jille/dfr v1.0.0
+	github.com/creachadair/getpass v0.2.1
 	github.com/google/go-cmp v0.5.8
 	github.com/iancoleman/strcase v0.2.0
 	github.com/jackc/pgx/v4 v4.16.1
@@ -27,5 +28,7 @@ require (
 	github.com/xdg-go/stringprep v1.0.3 // indirect
 	golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97 // indirect
 	golang.org/x/exp v0.0.0-20220613132600-b0d781184e0d // indirect
+	golang.org/x/sys v0.4.0 // indirect
+	golang.org/x/term v0.4.0 // indirect
 	golang.org/x/text v0.3.7 // indirect
 )
diff --git a/go.sum b/go.sum
index c9c55c2..192122f 100644
--- a/go.sum
+++ b/go.sum
@@ -8,6 +8,8 @@ github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
 github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
+github.com/creachadair/getpass v0.2.1 h1:ZtLVMUa5HdBStZsvSoYv4+G+RxQM1vNrfCUc0/xIJW8=
+github.com/creachadair/getpass v0.2.1/go.mod h1:nz1KzZI7LDzvc3d8CmJfGC1Hkiu6zw1iPuevAFiBJ+Y=
 github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7DoTY=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
@@ -169,8 +171,12 @@ golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.4.0 h1:Zr2JFtRQNX3BCZ8YtxRE9hNJYC8J6I1MVbMg6owUp18=
+golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.4.0 h1:O7UWfv5+A2qiuulQk30kVinPoMtoIPeVaKLEgLpVkvg=
+golang.org/x/term v0.4.0/go.mod h1:9P2UbLfCdcvo3p/nzKvsmas4TnlujnuoV9hGgYzW1lQ=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=