From 62c28005a8ee351d9faac5b5e9591edc48d10425 Mon Sep 17 00:00:00 2001
From: raul <raul@bulgariu.xyz>
Date: Wed, 3 Jul 2024 10:06:06 +0200
Subject: [PATCH] Use TLS encryption by default

---
 cmd/root.go       |  1 +
 cmd/server.go     |  1 -
 cmd/serverFunc.go | 51 ++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index c469774..1a28036 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -43,6 +43,7 @@ func initConfig() {
 
 		// Search config in home directory with name ".tune" (without extension).
 		viper.AddConfigPath(home + "/.config/tiamat")
+		configDir = home + "/.config/tiamat"
 		viper.SetConfigType("toml")
 		viper.SetConfigName("tiamat")
 	}
diff --git a/cmd/server.go b/cmd/server.go
index 45afdc5..caee177 100644
--- a/cmd/server.go
+++ b/cmd/server.go
@@ -9,7 +9,6 @@ import (
 	"log"
 )
 
-// serverCmd represents the server command
 var serverCmd = &cobra.Command{
 	Use:   "server",
 	Short: "Tiamat Server",
diff --git a/cmd/serverFunc.go b/cmd/serverFunc.go
index aa71a5b..0ed7ded 100644
--- a/cmd/serverFunc.go
+++ b/cmd/serverFunc.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"crypto/tls"
+	_ "embed"
 	"encoding/gob"
 	"encoding/json"
 	"fmt"
@@ -9,12 +11,16 @@ import (
 	"mime/multipart"
 	"net"
 	"os"
+	"os/exec"
+	"strings"
 	"time"
 
 	"github.com/spf13/viper"
 )
 
 var (
+	configDir string
+
 	C2Port     string = "1302"
 	clientList []Client
 	clientIDs  int = 0
@@ -22,11 +28,15 @@ var (
 	heartbeatRate time.Duration = 15
 
 	servInsecure bool
+	servCert     string
 
 	isUsingJSONParameter bool
 	clientJSONPath       string = "/.config/tiamat/clients.json"
 )
 
+//go:embed gen-cert.sh
+var script string
+
 func (c Client) Instruct(i Instructions) error {
 	enc := gob.NewEncoder(c.Conn)
 	err := enc.Encode(i)
@@ -36,6 +46,37 @@ func (c Client) Instruct(i Instructions) error {
 	return nil
 }
 
+func createCerts() {
+	log.Println("[-] Certificates don't exist! Creating them...")
+	c := exec.Command("bash")
+	c.Stdin = strings.NewReader(script)
+	b, err := c.Output()
+	if err != nil {
+		log.Fatalf("Error occurred creating certificates: %v\n", err)
+	}
+	fmt.Print(string(b))
+}
+
+func startInsecureServer() (net.Listener, error) {
+	ln, err := net.Listen("tcp", ":"+C2Port)
+	return ln, err
+}
+
+func startSecureServer() (net.Listener, error) {
+	cer, err := tls.LoadX509KeyPair(configDir+"/server.crt", configDir+"/server.key")
+
+	if os.IsNotExist(err) {
+		createCerts()
+		cer, err = tls.LoadX509KeyPair(configDir+"/server.crt", configDir+"/server.key")
+	}
+	if err != nil {
+		log.Fatalf("Error happened loading certificates: %v\n", err)
+	}
+	config := &tls.Config{Certificates: []tls.Certificate{cer}}
+	ln, err := tls.Listen("tcp", ":"+C2Port, config)
+	return ln, err
+}
+
 func setClientPath() (string, error) {
 	var fileToOpen string
 	home, err := os.UserHomeDir()
@@ -97,7 +138,15 @@ func Server() {
 		log.Fatalf("Error happened recovering clients: %v\n", err)
 	}
 
-	ln, err := net.Listen("tcp", ":"+C2Port)
+	//ln, err := net.Listen("tcp", ":"+C2Port)
+	var ln net.Listener
+	var err error
+	if servInsecure == true {
+		log.Println("WARNING: Starting unencrypted server!")
+		ln, err = startInsecureServer()
+	} else {
+		ln, err = startSecureServer()
+	}
 	if err != nil {
 		log.Fatalf("Error happened listening on C2 port: %v\n", err)
 	}