Merge pull request #12 from SlyngDK/request-ip

Support enroll with groups and requesting nebula ip

Author: SlyngDK

cmd/agent/enrollment.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/slackhq/nebula/cert"
@@ -31,8 +32,16 @@ var enrollCmd = &cobra.Command{
 		if err != nil {
 			l.WithError(err).Fatalln("failed to get token")
 		}
+		groups, err := cmd.Flags().GetString("groups")
+		if err != nil {
+			l.WithError(err).Fatalln("failed to get groups")
+		}
+		ip, err := cmd.Flags().GetString("ip")
+		if err != nil {
+			l.WithError(err).Fatalln("failed to get ip")
+		}
 
-		if err = enroll(agent, token); err != nil {
+		if err = enroll(agent, token, groups, ip); err != nil {
 			l.WithError(err).Fatalln("failed to enroll to server")
 		}
 	},
@@ -86,7 +95,16 @@ var enrollWaitCmd = &cobra.Command{
 		if status == 0 {
 			token, _ := cmd.Flags().GetString("token")
 			if token != "" {
-				if err := enroll(agent, token); err != nil {
+				groups, err := cmd.Flags().GetString("groups")
+				if err != nil {
+					l.WithError(err).Fatalln("failed to get groups")
+				}
+				ip, err := cmd.Flags().GetString("ip")
+				if err != nil {
+					l.WithError(err).Fatalln("failed to get ip")
+				}
+
+				if err := enroll(agent, token, groups, ip); err != nil {
 					l.WithError(err).Fatalln("failed to enroll to server")
 					os.Exit(2)
 				}
@@ -136,14 +154,19 @@ func getStatus(agent *agentClient) int8 {
 
 func init() {
 	enrollCmd.Flags().StringP("token", "t", "", "Enrollment token")
+	enrollCmd.Flags().StringP("groups", "g", "", "Comma separated list of groups")
+	enrollCmd.Flags().StringP("ip", "i", "", "Requesting for this specific nebula ip")
 	enrollCmd.MarkFlagRequired("token")
 	enrollWaitCmd.Flags().StringP("token", "t", "", "Enrollment token")
+	enrollWaitCmd.Flags().StringP("groups", "g", "", "Comma separated list of groups")
+	enrollWaitCmd.Flags().StringP("ip", "i", "", "Requesting for this specific nebula ip")
+	enrollWaitCmd.MarkFlagRequired("token")
 
 	enrollCmd.AddCommand(enrollStatusCmd)
 	enrollCmd.AddCommand(enrollWaitCmd)
 }
 
-func enroll(c *agentClient, enrollmentToken string) error {
+func enroll(c *agentClient, enrollmentToken, groups, ip string) error {
 	if enrollmentToken == "" {
 		return fmt.Errorf("requires enrollmentToken")
 	}
@@ -158,6 +181,17 @@ func enroll(c *agentClient, enrollmentToken string) error {
 		CsrPEM: string(csr),
 	}
 
+	if groups != "" {
+		g := strings.Split(groups, ",")
+		if len(g) > 0 {
+			enrollRequest.Groups = g
+		}
+	}
+
+	if ip != "" {
+		enrollRequest.RequestedIP = ip
+	}
+
 	hostname, err := os.Hostname()
 	if err == nil {
 		enrollRequest.Name = hostname

protocol/agent-service.pb.go

@@ -20,6 +20,7 @@ message EnrollRequest {
   string csrPEM = 2;
   repeated string groups = 3;
   string name = 4;
+  string requestedIP = 5;
 }
 
 message EnrollResponse {

protocol/agent-service.proto

@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"crypto/x509"
 	"fmt"
+	"net"
 	"strings"
 
 	"github.com/sirupsen/logrus"
@@ -58,6 +59,12 @@ func (a *agentService) Enroll(ctx context.Context, request *protocol.EnrollReque
 	if request.CsrPEM == "" {
 		return nil, status.Error(codes.InvalidArgument, "CsrPEM is required")
 	}
+	if request.RequestedIP != "" {
+		ip := net.ParseIP(request.RequestedIP)
+		if ip == nil {
+			return nil, status.Error(codes.InvalidArgument, "RequestedIP is not valid")
+		}
+	}
 
 	_, _, err = cert.UnmarshalX25519PublicKey([]byte(request.CsrPEM))
 	if err != nil {
@@ -68,7 +75,7 @@ func (a *agentService) Enroll(ctx context.Context, request *protocol.EnrollReque
 	addr := p.Addr.String()
 	ip := addr[0:strings.LastIndex(addr, ":")]
 
-	_, err = a.store.CreateEnrollmentRequest(fingerprint, request.Token, request.CsrPEM, ip, request.Name)
+	_, err = a.store.CreateEnrollmentRequest(fingerprint, request.Token, request.CsrPEM, ip, request.Name, request.RequestedIP, request.Groups)
 	if err != nil {
 		return nil, status.Error(codes.Internal, fmt.Sprintf("%s", err))
 	}