1 files changed, 52 insertions, 45 deletions

diff --git a/cmd/vhost/main.go b/cmd/vhost/main.go
index 2ecbfa7..f5016a1 100644
--- a/cmd/vhost/main.go
+++ b/cmd/vhost/main.go
@@ -24,72 +24,79 @@ func main() {
 
 	for scanner.Scan() {
 		line := scanner.Text()
-		
-		if line == "li" {
+		switch line {
+		case "li":
 			fmt.Println("Name\tAddr/Prefix\tState")
 			ipstack.SprintInterfaces()
-		}
-
-		if line == "ln" {
+		case "ln":
 			fmt.Println("Iface\tVIP\tUDPAddr")
 			ipstack.SprintNeighbors()
-		}
-
-		if line == "lr" {
+		case "lr":
 			fmt.Println("T\tPrefix\tNext Hop\tCost")
 			ipstack.SprintRoutingTable()
-		}
-		
-		if len(line) > 4 {
-			if line[:4] == "down" {
-				// get interface name
-				ifaceName := line[5:]
-				ipstack.InterfaceDownREPL(ifaceName)
-			}
+		case "q":
+			ipstack.CleanUp()
+			os.Exit(0)
+		default:
+			if len(line) > 4 {
+				if line[:4] == "down" {
+					ifaceName := line[5:]
+					ipstack.InterfaceDownREPL(ifaceName)
+				}
 
-			if line[:4] == "send" {
-				// get IP address and message that follows it
-				listOfWords := strings.Split(line, " ")
-				ipAddr := listOfWords[1]
-				message := listOfWords[2:]
-				// combine message into one string
-				messageToSend := strings.Join(message, " ")
-				// convert message to byte array
-				messageToSendBytes := []byte(messageToSend)
-				// get interface by ipAddr
-				iface, err := ipstack.GetNeighborByIP(ipAddr)
+				if line[:4] == "send" {
+					// get IP address and message that follows it
+					IPAndMessage := strings.Split(line, " ")
+					ipAddr := IPAndMessage[1]
+					message := IPAndMessage[2:]
+
+					// combine message into one string
+					messageToSend := strings.Join(message, " ")
+					messageToSendBytes := []byte(messageToSend)
+					
+					// check if ipAddr is in neighbor table
+					iface, err := ipstack.GetNeighborByIP(ipAddr)
 					if err != nil {
 						fmt.Println(err)
+
+						// check if ipAddr is in routing table
 						iface, err = ipstack.GetRouteByIP(ipAddr)
 						if err != nil {
 							fmt.Println(err)
 							continue
 						}
-						err = ipstack.SendIP(ipstack.GetMyVIP(), iface, 0, messageToSendBytes, ipAddr)
-						if err != nil {
-							fmt.Println(err)
-							continue
+
+						for _, interfaces := range ipstack.GetInterfaces() {
+							if interfaces.Name == iface.Name {
+								err = ipstack.SendIP(interfaces, iface, 0, messageToSendBytes, ipAddr)
+								if err != nil {
+									fmt.Println(err)
+								}
+								break
+							}
 						}
 						continue
 					}
-					err = ipstack.SendIP(ipstack.GetMyVIP(), iface, 0, messageToSendBytes, ipAddr)
-					if err != nil {
-						fmt.Println(err)
-						continue
+					for _, interfaces := range ipstack.GetInterfaces() {
+							if interfaces.Name == iface.Name {
+								err = ipstack.SendIP(interfaces, iface, 0, messageToSendBytes, ipAddr)
+								if err != nil {
+									fmt.Println(err)
+								}
+								break
+						}
 					}
+					continue
 				}
 			}
-		if len(line) > 2 {
-			if line[:2] == "up" {
-				// get interface name
-				ifaceName := line[3:]
-				ipstack.InterfaceUpREPL(ifaceName)
+			if len(line) > 2 {
+				if line[:2] == "up" {
+					// get interface name
+					ifaceName := line[3:]
+					ipstack.InterfaceUpREPL(ifaceName)
+				}
 			}
-		}
-
-		if line == "q" {
-			ipstack.CleanUp()
-			os.Exit(0)
+			continue
 		}
 	}
 }
\ No newline at end of file