package main

import (
	"bufio"
	"crypto/rand"
	"encoding/hex"
	flag "github.com/spf13/pflag"
	"fmt"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"io"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var workDir *string

func main() {
	// Use ConsoleWriter logger
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})

	// Create --send-to flag to send to a specific IP
	sendTo := flag.String("send-to", "", "Use IP address of receiver instead of mDNS")
	// Create --dest-dir flag to save to a specified folder
	destDir := flag.String("dest-dir", "", "Destination directory for files or dirs sent over opensend")
	workDir = flag.String("work-dir", "", "Working directory for opensend")
	givenCfgPath := flag.String("config", "", "Opensend config to use")
	// Create --skip-mdns to skip service registration
	skipMdns := flag.Bool("skip-mdns", false, "Skip zeroconf service registration (use if mdns fails)")
	// Create -t flag for type
	actionType := flag.StringP("type", "t", "", "Type of data being sent")
	// Create -d flag for data
	actionData := flag.StringP("data", "d", "", "Data to send")
	// Create -s flag for sending
	sendFlag := flag.BoolP("send", "s", false, "Send data")
	// Create -r flag for receiving
	recvFlag := flag.BoolP("receive", "r", false, "Receive data")
	targetFlag := flag.StringP("target", "T", "", "Target as defined in opensend.toml")
	// Parse flags
	flag.Parse()

	// Declare config variable
	var config *Config
	// If config flag not provided
	if *givenCfgPath == "" {
		// Get config path
		confPath := GetConfigPath()
		// Read config at path
		config = NewConfig(confPath)
	} else {
		// Otherwise, read config at provided path
		config = NewConfig(*givenCfgPath)
	}

	// If work directory flag not provided
	if *workDir == "" {
		// If send flag provided
		if *sendFlag {
			// Set work directory to sender as defined in config
			*workDir = ExpandPath(config.Sender.WorkDir)
		} else {
			// Otherwise set work directory to receiver as defined in config
			*workDir = ExpandPath(config.Receiver.WorkDir)
		}
	}

	// If target flag provided
	if *targetFlag != "" {
		// Set IP to target's IP
		*sendTo = config.Targets[*targetFlag].IP
	}

	// Create channel for signals
	sig := make(chan os.Signal, 1)
	// Send message on channel upon reception of SIGINT or SIGTERM
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	// Intercept signal
	go func() {
		select {
		// Wait for sig to be written
		case <-sig:
			// Warn user that a signal has been received and that opensend is shutting down
			log.Warn().Msg("Signal received. Shutting down.")
			// Remove opensend directory to avoid future conflicts
			_ = os.RemoveAll(*workDir)
			// Exit with code 0
			os.Exit(0)
		}
	}()

	// Create opensend dir ignoring errors
	_ = os.Mkdir(*workDir, 0755)
	// If -s given
	if *sendFlag {
		if *actionType == "" || *actionData == "" {
			log.Fatal().Msg("Valid action type and data is required to send")
		}
		// Create 32 byte buffer
		sharedKeyBytes := make([]byte, 32)
		// Read random bytes into buffer
		_, err := io.ReadFull(rand.Reader, sharedKeyBytes)
		if err != nil {
			log.Fatal().Err(err).Msg("Error generating random bytes")
		}
		// Encode random bytes to hexadecimal
		sharedKey := hex.EncodeToString(sharedKeyBytes)
		// Notify user a key has been created
		log.Info().Msg("Generated random shared key")
		// Create variable to store chosen IP
		var choiceIP string
		// If IP is provided via --send-to
		if *sendTo != "" {
			// Notify user that provided IP is being used
			log.Info().Msg("IP provided. Skipping discovery.")
			// Set chosen IP to provided
			choiceIP = *sendTo
			// Otherwise
		} else {
			// Notify user device discovery is beginning
			log.Info().Msg("Discovering opensend receivers")
			// Discover all _opensend._tcp.local. mDNS services
			discoveredReceivers, discoveredIPs := DiscoverReceivers()
			// Create reader for STDIN
			reader := bufio.NewReader(os.Stdin)
			// Print hostnames of each receiver
			for index, receiver := range discoveredReceivers {
				// Print hostname and index+1
				fmt.Println("["+strconv.Itoa(index+1)+"]", receiver)
			}
			// Prompt user for choice
			fmt.Print("Choose a receiver: ")
			choiceStr, _ := reader.ReadString('\n')
			// Convert input to int after trimming spaces
			choiceInt, err := strconv.Atoi(strings.TrimSpace(choiceStr))
			if err != nil {
				log.Fatal().Err(err).Msg("Error converting choice to int")
			}
			// Set choiceIndex to choiceInt-1 to allow for indexing
			choiceIndex := choiceInt - 1
			// Get IP of chosen receiver
			choiceIP = discoveredIPs[choiceIndex]
		}
		// Instantiate Config object
		parameters := NewParameters(*actionType, *actionData)
		// Validate data in config struct
		parameters.Validate()
		// Collect any files that may be required for transaction into opensend directory
		parameters.CollectFiles(*workDir)
		// Create config file in opensend directory
		parameters.CreateFile(*workDir)
		// Notify user of key exchange
		log.Info().Msg("Performing key exchange")
		// Exchange RSA keys with receiver
		rawKey := SenderKeyExchange(choiceIP)
		// Inform user receiver key has been received
		log.Info().Msg("Receiver key received")
		// Encrypt shared key using RSA public key
		key := EncryptKey(sharedKey, rawKey)
		// Save encrypted key in opensend directory as key.aes
		SaveEncryptedKey(key, *workDir+"/key.aes")
		// Notify user file encryption is beginning
		log.Info().Msg("Encrypting files")
		// Encrypt all files in opensend directory using shared key
		EncryptFiles(*workDir, sharedKey)
		// Notify user server has started
		log.Info().Msg("Server started on port 9898")
		// Send all files in opensend directory using an HTTP server on port 9898
		SendFiles(*workDir)
		// If -r given
	} else if *recvFlag {
		// If --skip-mdns is not given
		if !*skipMdns {
			// Register {hostname}._opensend._tcp.local. mDNS service and pass shutdown function
			zeroconfShutdown := RegisterService()
			// Shutdown zeroconf server at the end of main()
			defer zeroconfShutdown()
		}
		// Notify user keypair is being generated
		log.Info().Msg("Generating RSA keypair")
		// Generate keypair
		privateKey, publicKey := GenerateRSAKeypair()
		// Notify user opensend is waiting for key exchange
		log.Info().Msg("Waiting for sender key exchange")
		// Exchange keys with sender
		senderIP := ReceiverKeyExchange(publicKey)
		// Sleep 300ms to allow sender time to start HTTP server
		time.Sleep(300 * time.Millisecond)
		// Notify user files are being received
		log.Info().Msg("Receiving files from server (This may take a while)")
		// Connect to sender's TCP socket
		connection := ConnectToSender(senderIP)
		// Get files from sender and place them into the opensend directory
		RecvFiles(connection)
		// Get encrypted shared key from sender
		encryptedKey := GetKey(connection)
		// Send stop signal to sender's HTTP server
		SendSrvStopSignal(connection)
		// Decrypt shared key
		sharedKey := DecryptKey(encryptedKey, privateKey)
		// Notify user file decryption is beginning
		log.Info().Msg("Decrypting files")
		// Decrypt all files in opensend directory using shared key
		DecryptFiles(*workDir, sharedKey)
		// Instantiate Config
		parameters := &Parameters{}
		// Read config file in opensend directory
		parameters.ReadFile(*workDir + "/parameters.json")
		// Notify user that action is being executed
		log.Info().Msg("Executing JSON action")
		// Execute JSON action using files within opensend directory
		parameters.ExecuteAction(*workDir, *destDir)
	} else {
		flag.Usage()
		log.Fatal().Msg("You must choose sender or receiver mode using -s or -r")
	}
	// Remove opensend directory
	err := os.RemoveAll(*workDir)
	if err != nil {
		log.Fatal().Err(err).Msg("Error removing opensend directory")
	}
}