From f4becf1e73afc7ff6c12dd3228287acb0b8a4711 Mon Sep 17 00:00:00 2001
From: dovholuknf <46322585+dovholuknf@users.noreply.github.com>
Date: Fri, 4 Oct 2024 09:03:33 -0400
Subject: [PATCH 1/4] allow additional login params to be supplied

---
 .gitignore       |   3 +-
 zsshlib/flags.go |  18 +++----
 zsshlib/oidc.go  | 121 +++++++++++++++++++++++++++++++++++++++++++++--
 zsshlib/ssh.go   |  81 +------------------------------
 4 files changed, 131 insertions(+), 92 deletions(-)

diff --git a/.gitignore b/.gitignore
index c73ec0d..5f89bd0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,4 +4,5 @@ build
 zssh*json
 .run
 ziti-edge-tunnel*
-my.env
+*.env
+
diff --git a/zsshlib/flags.go b/zsshlib/flags.go
index 6ac15d8..8843db5 100644
--- a/zsshlib/flags.go
+++ b/zsshlib/flags.go
@@ -18,14 +18,15 @@ type SshFlags struct {
 }
 
 type OIDCFlags struct {
-	Mode          bool
-	Issuer        string
-	ClientID      string
-	ClientSecret  string
-	CallbackPort  string
-	AsAscii       bool
-	OIDCOnly      bool
-	ControllerUrl string
+	Mode                  bool
+	Issuer                string
+	ClientID              string
+	ClientSecret          string
+	CallbackPort          string
+	AsAscii               bool
+	OIDCOnly              bool
+	ControllerUrl         string
+	AdditionalLoginParams []string
 }
 
 type ScpFlags struct {
@@ -93,6 +94,7 @@ func (f *SshFlags) OIDCFlags(cmd *cobra.Command) {
 	cmd.Flags().BoolVarP(&f.OIDC.Mode, "oidc", "o", false, fmt.Sprintf("toggle OIDC mode. default: %t", defaults.OIDC.Enabled))
 	cmd.Flags().BoolVar(&f.OIDC.OIDCOnly, "oidcOnly", false, "toggle OIDC only mode. default: false")
 	cmd.Flags().StringVar(&f.OIDC.ControllerUrl, "controllerUrl", "", "the url of the controller to use. only used with --oidcOnly")
+	cmd.Flags().StringArrayVarP(&f.OIDC.AdditionalLoginParams, "additionalLoginParams", "l", []string{}, "Additional parameters to specify to the login. Can specify multiple times. Must be in the format of param=value")
 }
 
 func (f *SshFlags) AddCommonFlags(cmd *cobra.Command) {
diff --git a/zsshlib/oidc.go b/zsshlib/oidc.go
index 886c2be..3a6b1f1 100644
--- a/zsshlib/oidc.go
+++ b/zsshlib/oidc.go
@@ -2,9 +2,18 @@ package zsshlib
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"net/http"
+	"strings"
 	"time"
 
+	"github.com/google/uuid"
+	"github.com/sirupsen/logrus"
+	"github.com/zitadel/oidc/v2/pkg/client/rp"
+	"github.com/zitadel/oidc/v2/pkg/client/rp/cli"
+	httphelper "github.com/zitadel/oidc/v2/pkg/http"
+	"github.com/zitadel/oidc/v2/pkg/oidc"
 	"golang.org/x/oauth2"
 )
 
@@ -16,10 +25,11 @@ func OIDCFlow(initialContext context.Context, flags *SshFlags) (string, error) {
 			ClientSecret: flags.OIDC.ClientSecret,
 			RedirectURL:  fmt.Sprintf("http://localhost:%v%v", flags.OIDC.CallbackPort, callbackPath),
 		},
-		CallbackPath: callbackPath,
-		CallbackPort: flags.OIDC.CallbackPort,
-		Issuer:       flags.OIDC.Issuer,
-		Logf:         log.Debugf,
+		CallbackPath:          callbackPath,
+		CallbackPort:          flags.OIDC.CallbackPort,
+		Issuer:                flags.OIDC.Issuer,
+		Logf:                  log.Debugf,
+		AdditionalLoginParams: flags.OIDC.AdditionalLoginParams,
 	}
 	waitFor := 30 * time.Second
 	ctx, cancel := context.WithTimeout(initialContext, waitFor)
@@ -36,3 +46,106 @@ func OIDCFlow(initialContext context.Context, flags *SshFlags) (string, error) {
 
 	return token, nil
 }
+
+func zsshCodeFlow[C oidc.IDClaims](ctx context.Context, relyingParty rp.RelyingParty, config *OIDCConfig) *oidc.Tokens[C] {
+	codeflowCtx, codeflowCancel := context.WithCancel(ctx)
+	defer codeflowCancel()
+
+	tokenChan := make(chan *oidc.Tokens[C], 1)
+
+	callback := func(w http.ResponseWriter, r *http.Request, tokens *oidc.Tokens[C], state string, rp rp.RelyingParty) {
+		tokenChan <- tokens
+		msg := "<p><strong>Success!</strong></p>"
+		msg = msg + "<p>You are authenticated and can now return to the CLI.</p>"
+		w.Write([]byte(msg))
+	}
+
+	authHandlerWithQueryState := func(party rp.RelyingParty) http.HandlerFunc {
+		var urlParamOpts rp.URLParamOpt
+		for _, v := range config.AdditionalLoginParams {
+			parts := strings.Split(v, "=")
+			urlParamOpts = rp.WithURLParam(parts[0], parts[1])
+		}
+		return func(w http.ResponseWriter, r *http.Request) {
+			rp.AuthURLHandler(func() string {
+				return uuid.New().String()
+			}, party, urlParamOpts /*rp.WithURLParam("audience", "openziti2")*/)(w, r)
+		}
+	}
+
+	http.Handle("/login", authHandlerWithQueryState(relyingParty))
+	http.Handle(config.CallbackPath, rp.CodeExchangeHandler(callback, relyingParty))
+
+	httphelper.StartServer(codeflowCtx, ":"+config.CallbackPort)
+
+	cli.OpenBrowser("http://localhost:" + config.CallbackPort + "/login")
+
+	return <-tokenChan
+}
+
+// OIDCConfig represents a config for the OIDC auth flow.
+type OIDCConfig struct {
+	// CallbackPath is the path of the callback handler.
+	CallbackPath string
+
+	// CallbackPort is the port of the callback handler.
+	CallbackPort string
+
+	// Issuer is the URL of the OpenID Connect provider.
+	Issuer string
+
+	// HashKey is used to authenticate values using HMAC.
+	HashKey []byte
+
+	// BlockKey is used to encrypt values using AES.
+	BlockKey []byte
+
+	// IDToken is the ID token returned by the OIDC provider.
+	IDToken string
+
+	// Logger function for debug.
+	Logf func(format string, args ...interface{})
+
+	// Additional params to add to the login request
+	AdditionalLoginParams []string
+
+	oauth2.Config
+}
+
+// GetToken starts a local HTTP server, opens the web browser to initiate the OIDC Discovery and
+// Token Exchange flow, blocks until the user completes authentication and is redirected back, and returns
+// the OIDC tokens.
+func GetToken(ctx context.Context, config *OIDCConfig) (string, error) {
+	if err := config.validateAndSetDefaults(); err != nil {
+		return "", fmt.Errorf("invalid config: %w", err)
+	}
+
+	cookieHandler := httphelper.NewCookieHandler(config.HashKey, config.BlockKey, httphelper.WithUnsecure())
+
+	options := []rp.Option{
+		rp.WithCookieHandler(cookieHandler),
+		rp.WithVerifierOpts(rp.WithIssuedAtOffset(5 * time.Second)),
+	}
+	if config.ClientSecret == "" {
+		options = append(options, rp.WithPKCE(cookieHandler))
+	}
+
+	relyingParty, err := rp.NewRelyingPartyOIDC(config.Issuer, config.ClientID, config.ClientSecret, config.RedirectURL, config.Scopes, options...)
+	if err != nil {
+		logrus.Fatalf("error creating relyingParty %s", err.Error())
+	}
+
+	resultChan := make(chan *oidc.Tokens[*oidc.IDTokenClaims])
+
+	go func() {
+		tokens := zsshCodeFlow[*oidc.IDTokenClaims](ctx, relyingParty, config)
+		resultChan <- tokens
+	}()
+
+	select {
+	case tokens := <-resultChan:
+		return tokens.AccessToken, nil
+	case <-ctx.Done():
+		return "", errors.New("timeout: OIDC authentication took too long")
+	}
+}
diff --git a/zsshlib/ssh.go b/zsshlib/ssh.go
index e4bf4dc..ac50d4b 100644
--- a/zsshlib/ssh.go
+++ b/zsshlib/ssh.go
@@ -18,14 +18,8 @@ package zsshlib
 
 import (
 	"bufio"
-	"context"
 	"encoding/base64"
 	"fmt"
-	"github.com/google/uuid"
-	"github.com/gorilla/securecookie"
-	"github.com/zitadel/oidc/v2/pkg/client/rp/cli"
-	"github.com/zitadel/oidc/v2/pkg/oidc"
-	"golang.org/x/crypto/ssh/knownhosts"
 	"io"
 	"net"
 	"os"
@@ -36,15 +30,13 @@ import (
 	"sync"
 	"time"
 
+	"github.com/gorilla/securecookie"
 	"github.com/openziti/sdk-golang/ziti"
 	"github.com/pkg/errors"
 	"github.com/pkg/sftp"
-	"github.com/zitadel/oidc/v2/pkg/client/rp"
-	httphelper "github.com/zitadel/oidc/v2/pkg/http"
-	"golang.org/x/oauth2"
-
 	"github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/knownhosts"
 	"golang.org/x/crypto/ssh/terminal"
 )
 
@@ -140,75 +132,6 @@ func Dial(config *ssh.ClientConfig, conn net.Conn) (*ssh.Client, error) {
 	return ssh.NewClient(c, chans, reqs), nil
 }
 
-// OIDCConfig represents a config for the OIDC auth flow.
-type OIDCConfig struct {
-	// CallbackPath is the path of the callback handler.
-	CallbackPath string
-
-	// CallbackPort is the port of the callback handler.
-	CallbackPort string
-
-	// Issuer is the URL of the OpenID Connect provider.
-	Issuer string
-
-	// HashKey is used to authenticate values using HMAC.
-	HashKey []byte
-
-	// BlockKey is used to encrypt values using AES.
-	BlockKey []byte
-
-	// IDToken is the ID token returned by the OIDC provider.
-	IDToken string
-
-	// Logger function for debug.
-	Logf func(format string, args ...interface{})
-
-	oauth2.Config
-}
-
-// GetToken starts a local HTTP server, opens the web browser to initiate the OIDC Discovery and
-// Token Exchange flow, blocks until the user completes authentication and is redirected back, and returns
-// the OIDC tokens.
-func GetToken(ctx context.Context, config *OIDCConfig) (string, error) {
-	if err := config.validateAndSetDefaults(); err != nil {
-		return "", fmt.Errorf("invalid config: %w", err)
-	}
-
-	cookieHandler := httphelper.NewCookieHandler(config.HashKey, config.BlockKey, httphelper.WithUnsecure())
-
-	options := []rp.Option{
-		rp.WithCookieHandler(cookieHandler),
-		rp.WithVerifierOpts(rp.WithIssuedAtOffset(5 * time.Second)),
-	}
-	if config.ClientSecret == "" {
-		options = append(options, rp.WithPKCE(cookieHandler))
-	}
-
-	relyingParty, err := rp.NewRelyingPartyOIDC(config.Issuer, config.ClientID, config.ClientSecret, config.RedirectURL, config.Scopes, options...)
-	if err != nil {
-		logrus.Fatalf("error creating relyingParty %s", err.Error())
-	}
-
-	//ctx := context.Background()
-	state := func() string {
-		return uuid.New().String()
-	}
-
-	resultChan := make(chan *oidc.Tokens[*oidc.IDTokenClaims])
-
-	go func() {
-		tokens := cli.CodeFlow[*oidc.IDTokenClaims](ctx, relyingParty, config.CallbackPath, config.CallbackPort, state)
-		resultChan <- tokens
-	}()
-
-	select {
-	case tokens := <-resultChan:
-		return tokens.AccessToken, nil
-	case <-ctx.Done():
-		return "", errors.New("Timeout: OIDC authentication took too long")
-	}
-}
-
 // validateAndSetDefaults validates the config and sets default values.
 func (c *OIDCConfig) validateAndSetDefaults() error {
 	if c.ClientID == "" {

From 9ea7ac5f717d38e48e545790fedc713712bfc48c Mon Sep 17 00:00:00 2001
From: dovholuknf <46322585+dovholuknf@users.noreply.github.com>
Date: Fri, 4 Oct 2024 11:13:40 -0400
Subject: [PATCH 2/4] if supported, close the popped window...

---
 zsshlib/oidc.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zsshlib/oidc.go b/zsshlib/oidc.go
index 3a6b1f1..24966dc 100644
--- a/zsshlib/oidc.go
+++ b/zsshlib/oidc.go
@@ -55,7 +55,7 @@ func zsshCodeFlow[C oidc.IDClaims](ctx context.Context, relyingParty rp.RelyingP
 
 	callback := func(w http.ResponseWriter, r *http.Request, tokens *oidc.Tokens[C], state string, rp rp.RelyingParty) {
 		tokenChan <- tokens
-		msg := "<p><strong>Success!</strong></p>"
+		msg := "<script type=\"text/javascript\">window.close()</script><body onload=\"window.close()\">You may close this window</body><p><strong>Success!</strong></p>"
 		msg = msg + "<p>You are authenticated and can now return to the CLI.</p>"
 		w.Write([]byte(msg))
 	}

From 127ad808474e77c35000e61dfbd9253167d6ab10 Mon Sep 17 00:00:00 2001
From: dovholuknf <46322585+dovholuknf@users.noreply.github.com>
Date: Fri, 4 Oct 2024 11:53:43 -0400
Subject: [PATCH 3/4] fix -l missing

---
 zsshlib/oidc.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/zsshlib/oidc.go b/zsshlib/oidc.go
index 24966dc..044ed32 100644
--- a/zsshlib/oidc.go
+++ b/zsshlib/oidc.go
@@ -66,6 +66,11 @@ func zsshCodeFlow[C oidc.IDClaims](ctx context.Context, relyingParty rp.RelyingP
 			parts := strings.Split(v, "=")
 			urlParamOpts = rp.WithURLParam(parts[0], parts[1])
 		}
+		if urlParamOpts == nil {
+			urlParamOpts = func() []oauth2.AuthCodeOption {
+				return []oauth2.AuthCodeOption{}
+			}
+		}
 		return func(w http.ResponseWriter, r *http.Request) {
 			rp.AuthURLHandler(func() string {
 				return uuid.New().String()
@@ -144,6 +149,7 @@ func GetToken(ctx context.Context, config *OIDCConfig) (string, error) {
 
 	select {
 	case tokens := <-resultChan:
+		Logger().Infof("Refresh token: %s", tokens.RefreshToken)
 		return tokens.AccessToken, nil
 	case <-ctx.Done():
 		return "", errors.New("timeout: OIDC authentication took too long")

From 80124ec6f7116fbbacf3625028e61f7dc74239f3 Mon Sep 17 00:00:00 2001
From: dovholuknf <46322585+dovholuknf@users.noreply.github.com>
Date: Fri, 4 Oct 2024 13:10:09 -0400
Subject: [PATCH 4/4] update mm hook

---
 .github/workflows/mattermost-ziti-webhook.yml | 31 +++++++------------
 1 file changed, 12 insertions(+), 19 deletions(-)

diff --git a/.github/workflows/mattermost-ziti-webhook.yml b/.github/workflows/mattermost-ziti-webhook.yml
index f1ef9b4..5c4d942 100644
--- a/.github/workflows/mattermost-ziti-webhook.yml
+++ b/.github/workflows/mattermost-ziti-webhook.yml
@@ -1,4 +1,4 @@
-name: ziti-mattermost-action-py
+name: mattermost-ziti-webhook
 on:
   create:
   delete:
@@ -14,24 +14,17 @@ on:
   release:
     types: [released]
   workflow_dispatch:
-
 jobs:
-  ziti-webhook:
+  mattermost-ziti-webhook:
     runs-on: ubuntu-latest
-    name: Ziti Mattermost Action - Py
+    name: POST Webhook
+    if: github.repository_owner == 'openziti' && github.actor != 'dependabot[bot]'
+    env:
+      ZITI_LOG: 99
+      ZITI_NODEJS_LOG: 99
     steps:
-    - uses: openziti/ziti-mattermost-action-py@main
-      if: ${{ env.ZHOOK_URL != null }}
-      env:
-          ZHOOK_URL: ${{ secrets.ZHOOK_URL }}
-      with:
-        # Identity JSON containing key to access a Ziti network
-        zitiId: ${{ secrets.ZITI_MATTERMOST_IDENTITY }}
-
-        # URL to post the payload. Note that the `zitiId` must provide access to a service 
-        # intercepting `my-mattermost-ziti-server`
-        webhookUrl: ${{ secrets.ZHOOK_URL }}
-
-        eventJson: ${{ toJson(github.event) }}
-        senderUsername: "GitHubZ"
-        destChannel: "dev-notifications"
+      - uses: openziti/ziti-webhook-action@main
+        with:
+          ziti-id: ${{ secrets.ZITI_MATTERMOST_IDENTITY }}
+          webhook-url: ${{ secrets.ZITI_MATTERMOST_WEBHOOK_URL }}
+          webhook-secret: ${{ secrets.ZITI_MATTERMOSTI_WEBHOOK_SECRET }}
\ No newline at end of file