Optionize the notifier. Allow notifier to single-target commands (i.e…

…. allow using MicroMDM).
jessepeterson · May 1, 2023 · 186f5dd · 186f5dd
1 parent d3f961f
commit 186f5dd
Show file tree

Hide file tree

Showing 3 changed files with 129 additions and 52 deletions.
diff --git a/cmd/kmfddm/main.go b/cmd/kmfddm/main.go
@@ -44,6 +44,7 @@ func main() {
 		flEnqueueURL = flag.String("enqueue", "http://[::1]:9000/v1/enqueue/", "URL of NanoMDM enqueue endpoint")
 		flEnqueueKey = flag.String("enqueue-key", "", "NanoMDM API key")
 		flCORSOrigin = flag.String("cors-origin", "", "CORS Origin; for browser-based API access")
+		flMicro      = flag.Bool("micromdm", false, "Use MicroMDM command API calling conventions")
 	)
 	flag.Parse()
 
@@ -67,7 +68,13 @@ func main() {
 		os.Exit(1)
 	}
 
-	nanoNotif := notifier.New(storage, *flEnqueueURL, *flEnqueueKey, logger.With("service", "notifier"))
+	nOpts := []notifier.Option{
+		notifier.WithLogger(logger.With("service", "notifier")),
+	}
+	if *flMicro {
+		nOpts = append(nOpts, notifier.WithMicroMDM())
+	}
+	nanoNotif := notifier.New(storage, *flEnqueueURL, *flEnqueueKey, nOpts...)
 
 	mux := flow.New()
 

diff --git a/notifier/enqueue.go b/notifier/enqueue.go
@@ -20,63 +20,105 @@ func (n *Notifier) sendCommand(ctx context.Context, ids []string) error {
 		return errors.New("sending command: no ids")
 	}
 
-	logs := []interface{}{"msg", "sending command", "count", len(ids)}
-	defer func() { n.logger.Info(logs...) }()
-
-	var err error
-	var tokens []byte
-	if n.sendTokensForSingleID && len(ids) == 1 {
-		logs = append(logs, "include tokens", true)
-		tokens, err = n.store.RetrieveTokensJSON(ctx, ids[0])
-		if err != nil {
-			logs = append(logs, "err", err)
-			return fmt.Errorf("retrieving tokens json: %w", err)
+	var cmdIDs [][]string
+	if n.multi {
+		// multiple command targeting support
+		cmdIDs = [][]string{ids}
+	} else {
+		// split each id into a separate array
+		for _, id := range ids {
+			cmdIDs = append(cmdIDs, []string{id})
 		}
 	}
 
-	c := NewDeclarativeManagementCommand()
-	c.CommandUUID = uuid.NewString()
-	if len(tokens) > 0 {
-		c.Command.Data = &tokens
-	}
+	for _, ids := range cmdIDs {
+		if len(ids) < 1 {
+			continue
+		}
+		logs := []interface{}{
+			"msg", "sending command",
+			"count", len(ids),
+			"id_first", ids[0],
+		}
 
-	logs = append(logs, "command_uuid", c.CommandUUID)
+		var err error
+		var tokens []byte
+		if len(ids) == 1 && n.sendTokensForSingleID {
+			logs = append(logs, "include tokens", true)
+			tokens, err = n.store.RetrieveTokensJSON(ctx, ids[0])
+			if err != nil {
+				errLogs := append(logs,
+					"err",
+					fmt.Errorf("retrieving tokens json: %w", err),
+				)
+				n.logger.Info(errLogs...)
+				continue
+			}
+		}
 
-	cmdBytes, err := plist.Marshal(c)
-	if err != nil {
-		logs = append(logs, "err", err)
-		return err
-	}
+		c := NewDeclarativeManagementCommand()
+		c.CommandUUID = uuid.NewString()
+		if len(tokens) > 0 {
+			c.Command.Data = &tokens
+		}
 
-	var cmd io.Reader
-	if len(cmdBytes) > 0 {
-		cmd = bytes.NewBuffer(cmdBytes)
-	}
+		logs = append(logs, "command_uuid", c.CommandUUID)
 
-	// TODO: this probably won't scale well. the URL could proabably
-	// reasonably store 2K of data, so for a typical UUID this works out
-	// to 50-ish IDs in a single enqueuing.
-	idsStr := strings.Join(ids, ",")
-	req, err := http.NewRequestWithContext(ctx, http.MethodPut, n.url+idsStr, cmd)
-	if err != nil {
-		logs = append(logs, "err", err)
-		return err
-	}
-	req.SetBasicAuth("nanomdm", n.key)
-	req.Header.Set("User-Agent", agent)
+		cmdBytes, err := plist.Marshal(c)
+		if err != nil {
+			errLogs := append(logs,
+				"err",
+				fmt.Errorf("marshal command plist: %w", err),
+			)
+			n.logger.Info(errLogs...)
+			continue
+		}
 
-	resp, err := http.DefaultClient.Do(req)
-	if err != nil {
-		logs = append(logs, "err", err)
-		return err
-	}
-	defer resp.Body.Close()
-	logs = append(logs, "http_status", resp.StatusCode)
+		var cmd io.Reader
+		if len(cmdBytes) > 0 {
+			cmd = bytes.NewBuffer(cmdBytes)
+		}
+
+		// TODO: this probably won't scale well. the URL could proabably
+		// reasonably store 2K of data, so for a typical UUID this works out
+		// to 50-ish IDs in a single enqueuing.
+		idsStr := strings.Join(ids, ",")
+		req, err := http.NewRequestWithContext(ctx, http.MethodPut, n.url+idsStr, cmd)
+		if err != nil {
+			errLogs := append(logs,
+				"err",
+				fmt.Errorf("new request: %w", err),
+			)
+			n.logger.Info(errLogs...)
+			continue
+		}
+		req.SetBasicAuth(n.user, n.key)
+		req.Header.Set("User-Agent", agent)
+
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			errLogs := append(logs,
+				"err",
+				fmt.Errorf("http reqeust: %w", err),
+			)
+			n.logger.Info(errLogs...)
+			continue
+		}
+		resp.Body.Close()
+
+		logs = append(logs, "http_status", resp.StatusCode)
+
+		// TODO: read the success or failure of the command enqueing/pushing and report/error on it.
+		if resp.StatusCode != 200 && resp.StatusCode != 201 {
+			errLogs := append(logs,
+				"err",
+				fmt.Errorf("HTTP status: %s", resp.Status),
+			)
+			n.logger.Info(errLogs...)
+			continue
+		}
 
-	// TODO: read the success or failure of the command enqueing/pushing and report/error on it.
-	if resp.StatusCode != 200 {
-		logs = append(logs, "err", err)
-		return fmt.Errorf("HTTP status: %s", resp.Status)
+		n.logger.Debug(logs...)
 	}
 
 	return nil

diff --git a/notifier/notifier.go b/notifier/notifier.go
@@ -4,6 +4,7 @@ package notifier
 
 import (
 	"context"
+	"net/http"
 
 	"github.com/jessepeterson/kmfddm/log"
 )
@@ -18,20 +19,47 @@ type EnrollmentIDFinder interface {
 type Notifier struct {
 	store  EnrollmentIDFinder
 	url    string
+	user   string
+	method string
 	key    string
 	logger log.Logger
 
+	multi                 bool
 	sendTokensForSingleID bool
 }
 
-func New(store EnrollmentIDFinder, url, key string, logger log.Logger) *Notifier {
-	return &Notifier{
+type Option func(n *Notifier)
+
+func WithMicroMDM() Option {
+	return func(n *Notifier) {
+		n.user = "micromdm"
+		n.multi = false
+		n.method = http.MethodPost
+	}
+}
+
+func WithLogger(logger log.Logger) Option {
+	return func(n *Notifier) {
+		n.logger = logger
+	}
+}
+
+func New(store EnrollmentIDFinder, url, key string, opts ...Option) *Notifier {
+	n := &Notifier{
 		store:                 store,
 		url:                   url,
 		key:                   key,
-		logger:                logger,
+		logger:                log.NopLogger,
 		sendTokensForSingleID: true,
+
+		user:   "nanomdm",
+		method: http.MethodPut,
+		multi:  true,
+	}
+	for _, opt := range opts {
+		opt(n)
 	}
+	return n
 }
 
 func (n *Notifier) DeclarationChanged(ctx context.Context, declarationID string) error {