From cde920ce844879fe1f42d2e4b25d5dd5ab874ad5 Mon Sep 17 00:00:00 2001
From: Rustam Gilyazov <16064414+rusq@users.noreply.github.com>
Date: Sun, 7 Jan 2024 12:16:35 +1000
Subject: [PATCH] adding confirmation code into login flow

---
 auth/auth_ui/cli.go      |  6 +++++
 auth/auth_ui/huh.go      | 29 +++++++++++++++++++++++++
 auth/auth_ui/huh_test.go | 47 ++++++++++++++++++++++++++++++++++++++++
 auth/rod.go              | 18 +++++++++------
 4 files changed, 93 insertions(+), 7 deletions(-)
 create mode 100644 auth/auth_ui/huh_test.go

diff --git a/auth/auth_ui/cli.go b/auth/auth_ui/cli.go
index db8d1f75..b5b79799 100644
--- a/auth/auth_ui/cli.go
+++ b/auth/auth_ui/cli.go
@@ -109,3 +109,9 @@ func prompt(w io.Writer, prompt string, readlnFn func(*os.File) (string, error))
 		fmt.Fprintln(w, "input cannot be empty")
 	}
 }
+
+func (*CLI) ConfirmationCode(email string) (code int, err error) {
+	fmt.Printf("Enter confirmation code sent to %s: ", email)
+	_, err = fmt.Fscanf(os.Stdin, "%d", &code)
+	return
+}
diff --git a/auth/auth_ui/huh.go b/auth/auth_ui/huh.go
index cef4c9b5..f286ead4 100644
--- a/auth/auth_ui/huh.go
+++ b/auth/auth_ui/huh.go
@@ -4,6 +4,8 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"regexp"
+	"strconv"
 
 	"github.com/charmbracelet/huh"
 )
@@ -102,3 +104,30 @@ func valSepEaster() func(v int) error {
 		return nil
 	}
 }
+
+func (*Huh) ConfirmationCode(email string) (int, error) {
+	var strCode string
+	q := huh.NewInput().
+		CharLimit(6).
+		Title(fmt.Sprintf("Enter confirmation code sent to %s", email)).
+		Description("Slack did not recognise the browser, and sent a confirmation code.  Please enter the confirmation code below.").
+		Value(&strCode).
+		Validate(valSixDigits)
+	if err := q.Run(); err != nil {
+		return 0, err
+	}
+	code, err := strconv.Atoi(strCode)
+	if err != nil {
+		return 0, err
+	}
+	return code, nil
+}
+
+var numChlgRE = regexp.MustCompile(`^\d{6}$`)
+
+func valSixDigits(s string) error {
+	if numChlgRE.MatchString(s) {
+		return nil
+	}
+	return errors.New("confirmation code must be a sequence of six digits")
+}
diff --git a/auth/auth_ui/huh_test.go b/auth/auth_ui/huh_test.go
new file mode 100644
index 00000000..4e802a93
--- /dev/null
+++ b/auth/auth_ui/huh_test.go
@@ -0,0 +1,47 @@
+package auth_ui
+
+import "testing"
+
+func Test_valSixDigits(t *testing.T) {
+	type args struct {
+		s string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			"empty",
+			args{""},
+			true,
+		},
+		{
+			"too short",
+			args{"12345"},
+			true,
+		},
+		{
+			"too long",
+			args{"1234567"},
+			true,
+		},
+		{
+			"not a number",
+			args{"123456a"},
+			true,
+		},
+		{
+			"valid",
+			args{"123456"},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := valSixDigits(tt.args.s); (err != nil) != tt.wantErr {
+				t.Errorf("valSixDigits() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}
diff --git a/auth/rod.go b/auth/rod.go
index f313f200..596cdea2 100644
--- a/auth/rod.go
+++ b/auth/rod.go
@@ -7,7 +7,6 @@ import (
 	"os"
 	"time"
 
-	"github.com/charmbracelet/huh/spinner"
 	"github.com/rusq/slackauth"
 	"github.com/rusq/slackdump/v2/auth/auth_ui"
 )
@@ -29,6 +28,7 @@ type browserAuthUIExt interface {
 	BrowserAuthUI
 	RequestLoginType(w io.Writer) (int, error)
 	RequestCreds(w io.Writer, workspace string) (email string, passwd string, err error)
+	ConfirmationCode(email string) (code int, err error)
 }
 
 const expectedLoginDuration = 16 * time.Second
@@ -84,16 +84,20 @@ func NewRODAuth(ctx context.Context, opts ...Option) (RodAuth, error) {
 		if password == "" {
 			return r, fmt.Errorf("password cannot be empty")
 		}
+		fmt.Println("Logging in to Slack, depending on your connection speed, it usually takes 10-20 seconds...")
+
 		var loginErr error
-		spin := spinner.New().Title("Logging in...").Action(func() {
-			sp.Token, sp.Cookie, loginErr = slackauth.Headless(ctx, r.opts.workspace, username, password)
-		})
-		if err := spin.Run(); err != nil {
-			return r, err
-		}
+		sp.Token, sp.Cookie, loginErr = slackauth.Headless(
+			ctx,
+			r.opts.workspace,
+			username,
+			password,
+			slackauth.WithChallengeFunc(r.opts.ui.ConfirmationCode),
+		)
 		if loginErr != nil {
 			return r, loginErr
 		}
+
 		fmt.Fprintln(os.Stderr, "authenticated.")
 	case auth_ui.LoginCancel:
 		return r, ErrCancelled