From 00a9e1220c895efa43d37c9bfe65e163a89a0255 Mon Sep 17 00:00:00 2001
From: Martin Albert <martin.albert.187@gmail.com>
Date: Thu, 6 Jul 2023 16:44:55 +0200
Subject: [PATCH] Prepare CLI

---
 cmd/baton-demisto/config.go | 15 +++++++++++++++
 cmd/baton-demisto/main.go   | 34 ++++++++--------------------------
 2 files changed, 23 insertions(+), 26 deletions(-)

diff --git a/cmd/baton-demisto/config.go b/cmd/baton-demisto/config.go
index bdc0b65..896dd8e 100644
--- a/cmd/baton-demisto/config.go
+++ b/cmd/baton-demisto/config.go
@@ -2,16 +2,31 @@ package main
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/conductorone/baton-sdk/pkg/cli"
+	"github.com/spf13/cobra"
 )
 
 // config defines the external configuration required for the connector to run.
 type config struct {
 	cli.BaseConfig `mapstructure:",squash"` // Puts the base config options in the same place as the connector options
+
+	AccessToken string `mapstructure:"token"`
+	Unsafe      bool   `mapstructure:"unsafe"`
 }
 
 // validateConfig is run after the configuration is loaded, and should return an error if it isn't valid.
 func validateConfig(ctx context.Context, cfg *config) error {
+	if cfg.AccessToken == "" {
+		return fmt.Errorf("an access token must be provided")
+	}
+
 	return nil
 }
+
+// cmdFlags sets the cmdFlags required for the connector.
+func cmdFlags(cmd *cobra.Command) {
+	cmd.PersistentFlags().String("token", "", "Access token used to connect to the Cortex XSOAR API. ($BATON_TOKEN)")
+	cmd.PersistentFlags().Bool("unsafe", false, "Allow insecure TLS connections to Cortex XSOAR instance. ($BATON_UNSAFE)")
+}
diff --git a/cmd/baton-demisto/main.go b/cmd/baton-demisto/main.go
index a358206..cfe49af 100644
--- a/cmd/baton-demisto/main.go
+++ b/cmd/baton-demisto/main.go
@@ -5,8 +5,9 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/ConductorOne/baton-demisto/pkg/connector"
 	"github.com/conductorone/baton-sdk/pkg/cli"
-	"github.com/conductorone/baton-sdk/pkg/sdk"
+	"github.com/conductorone/baton-sdk/pkg/connectorbuilder"
 	"github.com/conductorone/baton-sdk/pkg/types"
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
 	"go.uber.org/zap"
@@ -18,13 +19,14 @@ func main() {
 	ctx := context.Background()
 
 	cfg := &config{}
-	cmd, err := cli.NewCmd(ctx, "baton-demisto", cfg, validateConfig, getConnector, run)
+	cmd, err := cli.NewCmd(ctx, "baton-demisto", cfg, validateConfig, getConnector)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err.Error())
 		os.Exit(1)
 	}
 
 	cmd.Version = version
+	cmdFlags(cmd)
 
 	err = cmd.Execute()
 	if err != nil {
@@ -36,37 +38,17 @@ func main() {
 func getConnector(ctx context.Context, cfg *config) (types.ConnectorServer, error) {
 	l := ctxzap.Extract(ctx)
 
-	c, err := sdk.NewEmptyConnector()
+	demistoConnector, err := connector.New(ctx, cfg.AccessToken, cfg.Unsafe)
 	if err != nil {
 		l.Error("error creating connector", zap.Error(err))
 		return nil, err
 	}
 
-	return c, nil
-}
-
-// run is where the process of syncing with the connector is implemented.
-func run(ctx context.Context, cfg *config) error {
-	l := ctxzap.Extract(ctx)
-
-	c, err := getConnector(ctx, cfg)
+	connector, err := connectorbuilder.NewConnector(ctx, demistoConnector)
 	if err != nil {
 		l.Error("error creating connector", zap.Error(err))
-		return err
-	}
-
-	r, err := sdk.NewConnectorRunner(ctx, c, cfg.C1zPath)
-	if err != nil {
-		l.Error("error creating connector runner", zap.Error(err))
-		return err
-	}
-	defer r.Close()
-
-	err = r.Run(ctx)
-	if err != nil {
-		l.Error("error running connector", zap.Error(err))
-		return err
+		return nil, err
 	}
 
-	return nil
+	return connector, nil
 }