From fb23d238a0c768e4b5f6c58ad417f28d8ad69071 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Mon, 16 Sep 2024 21:58:03 -0400
Subject: [PATCH] working 'zrok agent access private' (#463)

---
 agent/agent.go                 | 18 ++++++++++
 cmd/zrok/agentAccessPrivate.go | 65 ++++++++++++++++++++++++++++++++++
 cmd/zrok/agentSharePrivate.go  |  2 --
 cmd/zrok/main.go               |  6 ++++
 4 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 cmd/zrok/agentAccessPrivate.go

diff --git a/agent/agent.go b/agent/agent.go
index 72d9175e..7eb87f75 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -103,6 +103,24 @@ func (a *Agent) manager() {
 			} else {
 				logrus.Debug("skipping unidentified (orphaned) share removal")
 			}
+
+		case inAccess := <-a.inAccesses:
+			logrus.Infof("adding new access '%v'", inAccess.frontendToken)
+			a.accesses[inAccess.frontendToken] = inAccess
+
+		case outAccess := <-a.outAccesses:
+			if outAccess.frontendToken != "" {
+				logrus.Infof("removing access '%v'", outAccess.frontendToken)
+				if err := proctree.StopChild(outAccess.process); err != nil {
+					logrus.Errorf("error stopping access '%v': %v", outAccess.frontendToken, err)
+				}
+				if err := proctree.WaitChild(outAccess.process); err != nil {
+					logrus.Errorf("error joining access '%v': %v", outAccess.frontendToken, err)
+				}
+				delete(a.accesses, outAccess.frontendToken)
+			} else {
+				logrus.Debug("skipping unidentified (orphaned) access removal")
+			}
 		}
 	}
 }
diff --git a/cmd/zrok/agentAccessPrivate.go b/cmd/zrok/agentAccessPrivate.go
new file mode 100644
index 00000000..d3517ccc
--- /dev/null
+++ b/cmd/zrok/agentAccessPrivate.go
@@ -0,0 +1,65 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"github.com/openziti/zrok/agent/agentClient"
+	"github.com/openziti/zrok/agent/agentGrpc"
+	"github.com/openziti/zrok/environment"
+	"github.com/openziti/zrok/tui"
+	"github.com/spf13/cobra"
+)
+
+func init() {
+	agentAccessCmd.AddCommand(newAgentAccessPrivateCommand().cmd)
+}
+
+type agentAccessPrivateCommand struct {
+	bindAddress     string
+	responseHeaders []string
+	cmd             *cobra.Command
+}
+
+func newAgentAccessPrivateCommand() *agentAccessPrivateCommand {
+	cmd := &cobra.Command{
+		Use:   "private <token>",
+		Short: "Bind a private access in the zrok Agent",
+		Args:  cobra.ExactArgs(1),
+	}
+	command := &agentAccessPrivateCommand{cmd: cmd}
+	cmd.Flags().StringVarP(&command.bindAddress, "bind", "b", "127.0.0.1:9191", "The address to bind the private frontend")
+	cmd.Flags().StringArrayVar(&command.responseHeaders, "response-header", []string{}, "Add a response header ('key:value')")
+	cmd.Run = command.run
+	return command
+}
+
+func (cmd *agentAccessPrivateCommand) run(_ *cobra.Command, args []string) {
+	root, err := environment.LoadRoot()
+	if err != nil {
+		if !panicInstead {
+			tui.Error("unable to load environment", err)
+		}
+		panic(err)
+	}
+
+	if !root.IsEnabled() {
+		tui.Error("unable to load environment; did you 'zrok enable'?", nil)
+	}
+
+	client, conn, err := agentClient.NewClient(root)
+	if err != nil {
+		tui.Error("error connecting to agent", err)
+	}
+	defer conn.Close()
+
+	acc, err := client.PrivateAccess(context.Background(), &agentGrpc.PrivateAccessRequest{
+		Token:           args[0],
+		BindAddress:     cmd.bindAddress,
+		ResponseHeaders: cmd.responseHeaders,
+	})
+	if err != nil {
+		tui.Error("error creating access", err)
+	}
+
+	fmt.Println(acc)
+}
diff --git a/cmd/zrok/agentSharePrivate.go b/cmd/zrok/agentSharePrivate.go
index 8021c04f..eae79dea 100644
--- a/cmd/zrok/agentSharePrivate.go
+++ b/cmd/zrok/agentSharePrivate.go
@@ -19,7 +19,6 @@ func init() {
 
 type agentSharePrivateCommand struct {
 	backendMode  string
-	headless     bool
 	insecure     bool
 	closed       bool
 	accessGrants []string
@@ -34,7 +33,6 @@ func newAgentSharePrivateCommand() *agentSharePrivateCommand {
 	}
 	command := &agentSharePrivateCommand{cmd: cmd}
 	cmd.Flags().StringVarP(&command.backendMode, "backend-mode", "b", "proxy", "The backend mode {proxy, web, tcpTunnel, udpTunnel, caddy, drive, socks, vpn}")
-	cmd.Flags().BoolVar(&command.headless, "headless", false, "Disable TUI and run headless")
 	cmd.Flags().BoolVar(&command.insecure, "insecure", false, "Enable insecure TLS certificate validation for <target>")
 	cmd.Flags().BoolVar(&command.closed, "closed", false, "Enable closed permission mode (see --access-grant)")
 	cmd.Flags().StringArrayVar(&command.accessGrants, "access-grant", []string{}, "zrok accounts that are allowed to access this share (see --closed)")
diff --git a/cmd/zrok/main.go b/cmd/zrok/main.go
index 0a92da47..03dd946e 100644
--- a/cmd/zrok/main.go
+++ b/cmd/zrok/main.go
@@ -25,6 +25,7 @@ func init() {
 	adminCmd.AddCommand(adminListCmd)
 	adminCmd.AddCommand(adminUpdateCmd)
 	rootCmd.AddCommand(agentCmd)
+	agentCmd.AddCommand(agentAccessCmd)
 	agentCmd.AddCommand(agentShareCmd)
 	agentCmd.AddCommand(agentReleaseCmd)
 	testCmd.AddCommand(loopCmd)
@@ -80,6 +81,11 @@ var adminUpdateCmd = &cobra.Command{
 	Short: "Update global resources",
 }
 
+var agentAccessCmd = &cobra.Command{
+	Use:   "access",
+	Short: "zrok Agent access commands",
+}
+
 var agentCmd = &cobra.Command{
 	Use:     "agent",
 	Short:   "zrok Agent commands",