Add support to restart EAPOL and DHCP for a specific UNI

Change-Id: Idacb94e6a9e91221f2f008ec9b5d63210d74b7ce
diff --git a/internal/bbsim/api/onus_handler.go b/internal/bbsim/api/onus_handler.go
index 4ea60e7..91a107d 100644
--- a/internal/bbsim/api/onus_handler.go
+++ b/internal/bbsim/api/onus_handler.go
@@ -21,6 +21,7 @@
 	"fmt"
 	"github.com/opencord/bbsim/internal/bbsim/types"
 	"github.com/opencord/voltha-protos/v4/go/openolt"
+	"strconv"
 
 	"github.com/opencord/bbsim/api/bbsim"
 	"github.com/opencord/bbsim/internal/bbsim/devices"
@@ -341,19 +342,26 @@
 	return res, nil
 }
 
-func (s BBSimServer) RestartEapol(ctx context.Context, req *bbsim.ONURequest) (*bbsim.Response, error) {
-	// NOTE this API will change the EAPOL state for all UNIs on the requested ONU
-	// TODO a new API needs to be created to individually manage the UNIs
+func (s BBSimServer) RestartEapol(ctx context.Context, req *bbsim.UNIRequest) (*bbsim.Response, error) {
+	// NOTE this API will change the EAPOL state for all UNIs on the requested ONU if no UNI is specified
+	// Otherwise, it will change the EAPOL state for only the specified UNI on the ONU
 
 	res := &bbsim.Response{}
 
-	logger.WithFields(log.Fields{
-		"OnuSn": req.SerialNumber,
-	}).Infof("Received request to restart authentication ONU")
+	if req.UniID == "" {
+		logger.WithFields(log.Fields{
+			"OnuSn": req.OnuSerialNumber,
+		}).Infof("Received request to restart authentication for all UNIs on the ONU")
+	} else {
+		logger.WithFields(log.Fields{
+			"OnuSn": req.OnuSerialNumber,
+			"UniId": req.UniID,
+		}).Infof("Received request to restart authentication on UNI")
+	}
 
 	olt := devices.GetOLT()
 
-	onu, err := olt.FindOnuBySn(req.SerialNumber)
+	onu, err := olt.FindOnuBySn(req.OnuSerialNumber)
 
 	if err != nil {
 		res.StatusCode = int32(codes.NotFound)
@@ -365,8 +373,13 @@
 	startedOn := []string{}
 	success := true
 
+	uniIDint, err := strconv.Atoi(req.UniID)
 	for _, u := range onu.UniPorts {
 		uni := u.(*devices.UniPort)
+		//if a specific uni is specified, only restart it
+		if err == nil && req.UniID != "" && uni.ID != uint32(uniIDint) {
+			continue
+		}
 		if !uni.OperState.Is(devices.UniStateUp) {
 			// if the UNI is disabled, ignore it
 			continue
@@ -410,35 +423,65 @@
 			res.Message = "No service requires EAPOL"
 		}
 
-		logger.WithFields(log.Fields{
-			"OnuSn":   req.SerialNumber,
-			"Message": res.Message,
-		}).Info("Processed EAPOL restart request for ONU")
+		if req.UniID == "" {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+			}).Info("Processed EAPOL restart request for all UNIs on the ONU")
+		} else {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+				"UniId":   req.UniID,
+			}).Info("Processed EAPOL restart request on UNI")
+		}
+
 	} else {
 		res.StatusCode = int32(codes.FailedPrecondition)
 		res.Message = fmt.Sprintf("%v", errors)
 		logger.WithFields(log.Fields{
-			"OnuSn":   req.SerialNumber,
+			"OnuSn":   req.OnuSerialNumber,
 			"Message": res.Message,
-		}).Error("Error while processing DHCP restart request for ONU")
+		}).Error("Error while processing EAPOL restart request for ONU")
+
+		if req.UniID == "" {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+			}).Error("Error while processing EAPOL restart request for all UNIs on the ONU")
+		} else {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+				"UniId":   req.UniID,
+			}).Error("Error while processing EAPOL restart request on UNI")
+		}
+
 	}
 
 	return res, nil
 }
 
-func (s BBSimServer) RestartDhcp(ctx context.Context, req *bbsim.ONURequest) (*bbsim.Response, error) {
-	// NOTE this API will change the DHCP state for all UNIs on the requested ONU
-	// TODO a new API needs to be created to individually manage the UNIs
+func (s BBSimServer) RestartDhcp(ctx context.Context, req *bbsim.UNIRequest) (*bbsim.Response, error) {
+	// NOTE this API will change the DHCP state for all UNIs on the requested ONU if no UNI is specified
+	// Otherwise, it will change the DHCP state for only the specified UNI on the ONU
 
 	res := &bbsim.Response{}
 
-	logger.WithFields(log.Fields{
-		"OnuSn": req.SerialNumber,
-	}).Infof("Received request to restart DHCP on ONU")
+	if req.UniID == "" {
+		logger.WithFields(log.Fields{
+			"OnuSn": req.OnuSerialNumber,
+		}).Infof("Received request to restart authentication for all UNIs on the ONU")
+	} else {
+		logger.WithFields(log.Fields{
+			"OnuSn": req.OnuSerialNumber,
+			"UniId": req.UniID,
+		}).Infof("Received request to restart authentication on UNI")
+	}
 
 	olt := devices.GetOLT()
 
-	onu, err := olt.FindOnuBySn(req.SerialNumber)
+	onu, err := olt.FindOnuBySn(req.OnuSerialNumber)
 
 	if err != nil {
 		res.StatusCode = int32(codes.NotFound)
@@ -450,8 +493,13 @@
 	startedOn := []string{}
 	success := true
 
+	uniIDint, err := strconv.Atoi(req.UniID)
 	for _, u := range onu.UniPorts {
 		uni := u.(*devices.UniPort)
+		//if a specific uni is specified, only restart it
+		if err == nil && req.UniID != "" && uni.ID != uint32(uniIDint) {
+			continue
+		}
 		if !uni.OperState.Is(devices.UniStateUp) {
 			// if the UNI is disabled, ignore it
 			continue
@@ -460,7 +508,6 @@
 			service := s.(*devices.Service)
 			serviceKey := fmt.Sprintf("uni[%d]%s", uni.ID, service.Name)
 			if service.NeedsDhcp {
-
 				if err := service.DHCPState.Event("start_dhcp"); err != nil {
 					logger.WithFields(log.Fields{
 						"OnuId":   onu.ID,
@@ -486,17 +533,35 @@
 		} else {
 			res.Message = "No service requires DHCP"
 		}
-		logger.WithFields(log.Fields{
-			"OnuSn":   req.SerialNumber,
-			"Message": res.Message,
-		}).Info("Processed DHCP restart request for ONU")
+
+		if req.UniID == "" {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+			}).Info("Processed DHCP restart request for all UNIs on the ONU")
+		} else {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+				"UniId":   req.UniID,
+			}).Info("Processed DHCP restart request on UNI")
+		}
 	} else {
 		res.StatusCode = int32(codes.FailedPrecondition)
 		res.Message = fmt.Sprintf("%v", errors)
-		logger.WithFields(log.Fields{
-			"OnuSn":   req.SerialNumber,
-			"Message": res.Message,
-		}).Error("Error while processing DHCP restart request for ONU")
+
+		if req.UniID == "" {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+			}).Error("Error while processing DHCP restart request for all UNIs on the ONU")
+		} else {
+			logger.WithFields(log.Fields{
+				"OnuSn":   req.OnuSerialNumber,
+				"Message": res.Message,
+				"UniId":   req.UniID,
+			}).Error("Error while processing DHCP restart request on UNI")
+		}
 	}
 
 	return res, nil
diff --git a/internal/bbsimctl/commands/onu.go b/internal/bbsimctl/commands/onu.go
index 4ae8841..f929f39 100644
--- a/internal/bbsimctl/commands/onu.go
+++ b/internal/bbsimctl/commands/onu.go
@@ -245,8 +245,9 @@
 
 	ctx, cancel := context.WithTimeout(context.Background(), config.GlobalConfig.Grpc.Timeout)
 	defer cancel()
-	req := pb.ONURequest{
-		SerialNumber: string(options.Args.OnuSn),
+	req := pb.UNIRequest{
+		OnuSerialNumber: string(options.Args.OnuSn),
+		UniID:           "",
 	}
 	res, err := client.RestartEapol(ctx, &req)
 
@@ -266,8 +267,9 @@
 
 	ctx, cancel := context.WithTimeout(context.Background(), config.GlobalConfig.Grpc.Timeout)
 	defer cancel()
-	req := pb.ONURequest{
-		SerialNumber: string(options.Args.OnuSn),
+	req := pb.UNIRequest{
+		OnuSerialNumber: string(options.Args.OnuSn),
+		UniID:           "",
 	}
 	res, err := client.RestartDhcp(ctx, &req)
 
diff --git a/internal/bbsimctl/commands/uni.go b/internal/bbsimctl/commands/uni.go
index 81f7c96..b025e72 100644
--- a/internal/bbsimctl/commands/uni.go
+++ b/internal/bbsimctl/commands/uni.go
@@ -19,6 +19,7 @@
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"strconv"
 	"strings"
@@ -56,10 +57,26 @@
 	} `positional-args:"yes"`
 }
 
+type UNIEapolRestart struct {
+	Args struct {
+		OnuSn OnuSnString
+		UniId UniIdInt
+	} `positional-args:"yes" required:"yes"`
+}
+
+type UNIDhcpRestart struct {
+	Args struct {
+		OnuSn OnuSnString
+		UniId UniIdInt
+	} `positional-args:"yes" required:"yes"`
+}
+
 type UNIOptions struct {
-	List     UNIList     `command:"list"`
-	Get      UNIGet      `command:"get"`
-	Services UNIServices `command:"services"`
+	List         UNIList         `command:"list"`
+	Get          UNIGet          `command:"get"`
+	Services     UNIServices     `command:"services"`
+	RestartEapol UNIEapolRestart `command:"auth_restart"`
+	RestartDchp  UNIDhcpRestart  `command:"dhcp_restart"`
 }
 
 func RegisterUNICommands(parser *flags.Parser) {
@@ -142,6 +159,52 @@
 	return nil
 }
 
+func (options *UNIEapolRestart) Execute(args []string) error {
+	client, conn := connect()
+	defer conn.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), config.GlobalConfig.Grpc.Timeout)
+	defer cancel()
+
+	req := pb.UNIRequest{
+		OnuSerialNumber: string(options.Args.OnuSn),
+		UniID:           string(options.Args.UniId),
+	}
+	res, err := client.RestartEapol(ctx, &req)
+
+	if err != nil {
+		log.Fatalf("Cannot restart EAPOL for ONU %s and UNI %s: %v", options.Args.OnuSn, options.Args.UniId, err)
+		return err
+	}
+
+	fmt.Println(fmt.Sprintf("[Status: %d] %s", res.StatusCode, res.Message))
+
+	return nil
+}
+
+func (options *UNIDhcpRestart) Execute(args []string) error {
+	client, conn := connect()
+	defer conn.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), config.GlobalConfig.Grpc.Timeout)
+	defer cancel()
+
+	req := pb.UNIRequest{
+		OnuSerialNumber: string(options.Args.OnuSn),
+		UniID:           string(options.Args.UniId),
+	}
+	res, err := client.RestartDhcp(ctx, &req)
+
+	if err != nil {
+		log.Fatalf("Cannot restart DHCP for ONU %s and UNI %s: %v", options.Args.OnuSn, options.Args.UniId, err)
+		return err
+	}
+
+	fmt.Println(fmt.Sprintf("[Status: %d] %s", res.StatusCode, res.Message))
+
+	return nil
+}
+
 func (uniId *UniIdInt) Complete(match string) []flags.Completion {
 	client, conn := connect()
 	defer conn.Close()