package main

import (
    "crypto/sha256"
    "crypto/rand"
    "golang.org/x/crypto/chacha20poly1305"
    "github.com/atotto/clipboard"
    "encoding/base64"
    "encoding/gob"
    "encoding/hex"
    "fmt"
    "net"
    "time"
    "golang.org/x/term"
    "syscall"
    "flag"
    "os"
    "log"
    "io/ioutil"
    "gopkg.in/yaml.v2"
)

var ChaKey = []byte("")
type Config struct {
	Auth struct {
		ChaKey string `yaml:"ChaKey"`
	} `yaml:"auth"`
}

type Request struct {
    KeyPwd    string
    Username  string
    Password  string
    Operation string
    Site      string
    VaultData string
    Timestamp int64
    Nonce     string
}

type Response struct {
    Message string
    Enc     string
}

func generateNonce() (string, error) {
    b := make([]byte, 12) // 96-bit nonce
    _, err := rand.Read(b)
    if err != nil {
        return "", err
    }
    return hex.EncodeToString(b), nil
}

func getKey()([]byte) {
    // Hash it to 32 bytes using SHA-256
    hashedKey := sha256.Sum256(ChaKey)
    theKey := hashedKey[:] // Convert [32]byte to []byte
    return theKey
}

func chEnc(p string)(string) {
    // Create cipher
    aead, err := chacha20poly1305.NewX(getKey())
    if err != nil {
        log.Fatalf("ChaKey enc error")
        return ""
    }    
    // Create nonce
    nonce := make([]byte, aead.NonceSize())
    if _, err := rand.Read(nonce); err != nil {
        log.Fatalf("Cha Nonce Error")
        return ""
    }    
    // Encrypt the encoded data
    encrypted := aead.Seal(nil, nonce, []byte(p), nil)
        
    // Send nonce + encrypted data
    fullMessage := append(nonce, encrypted...)    
    encoded := base64.StdEncoding.EncodeToString(fullMessage)
    return encoded  
}

func chDec(eText string)(string) {
    if eText == "" {
        return ""
    }
    
    // Decoding from base64
	decoded, err := base64.StdEncoding.DecodeString(eText)
	if err != nil {
		log.Fatal("Base64 decoding error:", err)
	}

	// Create cipher instance (XChaCha20-Poly1305 for longer nonces)
	aead, err := chacha20poly1305.NewX(getKey())
	if err != nil {
        log.Fatalf("Client: ChaKey failed")
		return ""
	}

    encryptedMsg := decoded

	// Decrypt: Split nonce and ciphertext
	decryptedNonce := encryptedMsg[:aead.NonceSize()]
	decryptedCiphertext := encryptedMsg[aead.NonceSize():]

	decrypted, err := aead.Open(nil, decryptedNonce, decryptedCiphertext, nil)
	if err != nil {
        log.Fatalf("Client decryption failed")
		return ""
	}
    
    return string(decrypted)
}

func copyText(text string) {
    err := clipboard.WriteAll(text)
	if err != nil {
		fmt.Println("Failed to copy to clipboard:", err)
		return
	}
	fmt.Println("Text copied to clipboard!")
}

func pause() {
    var input string
	fmt.Print("Press Enter to continue...")
	fmt.Scanln(&input)
}

func main() {
    configFilePtr := flag.String("config", "config.yaml", "Path to the YAML configuration file")
	flag.Parse()
    
    if *configFilePtr == "" {
		fmt.Println("Please specify a configuration file using -config")
		flag.Usage()
        pause()
		os.Exit(1)
	}

	// Read the config file
	yamlFile, err := ioutil.ReadFile(*configFilePtr)
	if err != nil {
		log.Fatalf("Error reading YAML file: %v\n", err)
        pause()
	}

	// Parse the YAML
	var config Config
	err = yaml.Unmarshal(yamlFile, &config)
	if err != nil {
		log.Fatalf("Error parsing YAML file: %v\n", err)
        pause()
	}

	// Print the ChaKey
	if config.Auth.ChaKey == "" {
		fmt.Println("Warning: ChaKey not found in configuration file")
        pause()
    }
    ChaKey = []byte(config.Auth.ChaKey)
    
    var host string
    fmt.Print("Enter host or IP: ")
    fmt.Scanln(&host)

    conn, err := net.Dial("tcp", host+":9898")
    if err != nil {
        panic(err)
    }
    defer conn.Close()
    
    fmt.Print("Key Passphrase: ")
    byteKeyPwd, err := term.ReadPassword(int(syscall.Stdin))
    fmt.Println()
    if err != nil {
        panic(err)
    }
    keypassphrase := string(byteKeyPwd)

    var username, op, site, vaultData string
    fmt.Print("Username: ")
    fmt.Scanln(&username)
    
    fmt.Print("Password: ")
    bytePassword, err := term.ReadPassword(int(syscall.Stdin))
    fmt.Println()
    if err != nil {
        panic(err)
    }
    password := string(bytePassword)

    fmt.Print("Operation (register/store/get): ")
    fmt.Scanln(&op)

    if op == "store" || op == "get" {
        fmt.Print("Site: ")
        fmt.Scanln(&site)
    }
    if op == "store" {
        fmt.Print("Password to store: ")
        fmt.Scanln(&vaultData)
    }

    nonce, err := generateNonce()
    if err != nil {
        panic(err)
    }

    var myPwd string
    if vaultData == "" {
        myPwd = ""
    } else {
        myPwd = chEnc(vaultData)
    }

    req := Request{
        KeyPwd:    chEnc(keypassphrase), 
        Username:  chEnc(username),
        Password:  chEnc(password),
        Operation: op,
        Site:      chEnc(site),
        VaultData: myPwd,
        Timestamp: time.Now().Unix(),
        Nonce:     nonce,
    }

    enc := gob.NewEncoder(conn)
    dec := gob.NewDecoder(conn)

    if err := enc.Encode(req); err != nil {
        fmt.Println("Failed to send request:", err)
        pause()
        return
    }

    var res Response
    if err := dec.Decode(&res); err != nil {
        fmt.Println("Failed to receive response:", err)
        pause()
        return
    }

    if res.Message != "" {
        fmt.Println("Server response:", res.Message)
    } else {
        passwd := chDec(res.Enc)
        fmt.Println("Server response:", passwd)
        copyText(passwd)
    }
    pause()
}