// Package cli defines flags, commands and functions to run.
package cli

import (
	"bytes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/AlecAivazis/survey/v2"
	"github.com/BurntSushi/toml"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"github.com/urfave/cli/v2"
	"github.com/urfave/cli/v2/altsrc"
	entropy "github.com/wagslane/go-password-validator"
	"golang.org/x/crypto/argon2"
	"golang.org/x/crypto/chacha20poly1305"
	"net"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

// config represents the configuration file
type config struct {
	Address string        `toml:"address"`
	Port    []string      `toml:"port"`
	Timeout time.Duration `toml:"timeout"`
	Delay   time.Duration `toml:"delay"`
}

// ding holds variables to operate
type ding struct {
	cliCtx *cli.Context
	config
	configFile *os.File
}

const (
	appName = "ding"
)

var (
	// version and commit can be defined at compilation time via ldflags.
	version, commit string
)

func Run() error {
	// Use all processors
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Set global logger
	log.Logger = zerolog.New(zerolog.ConsoleWriter{Out: os.Stdout, TimeFormat: time.RFC3339}).With().Timestamp().Logger()

	// Common flag
	var noCipher = &cli.BoolFlag{Name: "insecure", Aliases: []string{"i"}, Usage: "don't de/cipher configuration file", Value: false}

	// Initialize ding
	app := cli.NewApp()

	globalFlags := []cli.Flag{
		altsrc.NewStringFlag(&cli.StringFlag{Name: "address", Aliases: []string{"a"}, Usage: "address to knock", Hidden: true}),
		altsrc.NewStringSliceFlag(&cli.StringSliceFlag{Name: "port", Aliases: []string{"p"}, Usage: "ports to knock", Hidden: true, Action: func(_ *cli.Context, ports []string) error { return validatePortRange(ports) }}),
		altsrc.NewDurationFlag(&cli.DurationFlag{Name: "timeout", Aliases: []string{"t"}, Usage: "timeout in milliseconds", Hidden: true, DefaultText: "1500ms", Value: 1500 * time.Millisecond}),
		altsrc.NewDurationFlag(&cli.DurationFlag{Name: "delay", Aliases: []string{"d"}, Usage: "delay in milliseconds between knocks", Hidden: true, DefaultText: "100ms", Value: 100 * time.Millisecond}),
		noCipher,
		&cli.StringFlag{Name: "config", Aliases: []string{"c"}, Usage: "input configuration file", Hidden: true, DefaultText: "$XDG_CONFIG_HOME/ding/config.toml or $HOME/.config/ding/config.toml"},
	}

	app.Name = appName
	if version == "" || commit == "" {
		version = "untagged"
		commit = "0000000000"
	}
	app.Version = version + "-" + commit[0:10]
	app.Usage = "Command line interface tool to knock ports"
	app.UseShortOptionHandling = true
	app.Suggest = true
	app.EnableBashCompletion = true
	app.Authors = []*cli.Author{
		{
			Name:  "Adrien",
			Email: "contact@illuad.fr",
		},
	}
	app.Flags = globalFlags
	app.Commands = cli.Commands{
		{
			Name:    "setup",
			Aliases: []string{"s"},
			Usage:   "Launches ding setup",
			Flags: []cli.Flag{
				&cli.StringFlag{Name: "address", Aliases: []string{"a"}, Usage: "address to knock"},
				&cli.StringSliceFlag{Name: "port", Aliases: []string{"p"}, Usage: "ports to knock", Action: func(_ *cli.Context, ports []string) error { return validatePortRange(ports) }},
				&cli.DurationFlag{Name: "timeout", Aliases: []string{"t"}, Usage: "timeout in milliseconds", DefaultText: "1500ms", Value: 1500 * time.Millisecond},
				&cli.DurationFlag{Name: "delay", Aliases: []string{"d"}, Usage: "delay in milliseconds between knocks", DefaultText: "100ms", Value: 100 * time.Millisecond},
				noCipher,
				&cli.BoolFlag{Name: "bypass-password-entropy", Aliases: []string{"b"}, Usage: "insecurely bypass password entropy", Value: false}},
			Action: func(cliCtx *cli.Context) error {
				dingCLI := ding{cliCtx: cliCtx}

				dingCLI.
					askAddressIfNotSet().
					askPortIfNotSet().
					askTimeoutIfNotSet().
					askDelayIfNotSet()

				if cliCtx.Bool("insecure") {
					return dingCLI.write()
				}

				return dingCLI.marshal()
			},
		},
	}
	app.Before = func(cliCtx *cli.Context) error {
		if cliCtx.IsSet("config") {
			return altsrc.InitInputSourceWithContext(cliCtx.App.Flags, altsrc.NewTomlSourceFromFlagFunc("config"))(cliCtx)
		}

		return nil
	}
	app.Action = func(cliCtx *cli.Context) error {
		dingCLI := ding{cliCtx: cliCtx}

		if cliCtx.Bool("insecure") {
			if err := dingCLI.read(); err != nil {
				return err
			}
		} else {
			if err := dingCLI.unmarshal(); err != nil {
				return err
			}
		}

		address := dingCLI.config.Address

		if _, err := net.LookupHost(address); err != nil {
			return err
		}

		for _, p := range dingCLI.config.Port {
			_, _ = net.DialTimeout("tcp", fmt.Sprintf("%s:%s", address, p), dingCLI.config.Timeout)

			time.Sleep(dingCLI.config.Delay)
		}

		return nil
	}

	return app.Run(os.Args)
}

// basePath returns configuration path, e.g. /home/jdoe/.config/ding.
// Configuration filename (config.toml) is not added, see configPath for this.
// Empty string is returned if path can't be build (if $XDG_CONFIG_HOME or $HOME are not set).
func (d *ding) basePath() string {
	path, ok := os.LookupEnv("XDG_CONFIG_HOME")
	if !ok {
		// $XDG_CONFIG_HOME is not set
		// $HOME is?
		path, ok = os.LookupEnv("HOME")
		if !ok {
			// $HOME is not set
			// Base path can't be build
			return ""
		} else {
			// $HOME is set, add ".config/ding"
			return filepath.Join(path, ".config", "ding")
		}
	} else {
		// $XDG_CONFIG_HOME is set, add "ding"
		return filepath.Join(path, "ding")
	}
}

// configPath return configuration file path, e.g. /home/jdoe/.config/ding/config.toml.
// Empty string is returned if configuration file path can't be build (if basePath has returned empty string).
func (d *ding) configPath() string {
	path := d.basePath()
	if path == "" {
		return ""
	}

	return filepath.Join(path, "config.toml")
}

// saltPath return salt file path, e.g. /home/jdoe/.config/ding/.salt.
// Empty string is returned if salt file path can't be build (if basePath has returned empty string).
func (d *ding) saltPath() string {
	path := d.basePath()
	if path == "" {
		return ""
	}

	return filepath.Join(path, ".salt")
}

// createBaseFolders creates configuration path, e.g. /home/jdoe/.config/ding.
// Configuration file (config.toml) is not created, see createConfigurationFile for this.
// An error is returned if basePath has returned empty string.
func (d *ding) createBaseFolders() (string, error) {
	path := d.basePath()
	if path == "" {
		return "", errors.New("failed to retrieve base path")
	}

	return path, os.MkdirAll(path, 0755)
}

// createConfigurationFile creates configuration file, e.g. /home/jdoe/.config/ding/config.toml.
// An error is return if createBaseFolders failed to create base folders or if configuration file creation failed.
func (d *ding) createConfigurationFile() error {
	path, err := d.createBaseFolders()
	if err != nil {
		return err
	}

	d.configFile, err = os.Create(filepath.Join(path, "config.toml"))

	return err
}

// askAddressIfNotSet asks for the address, or sets it if given by flag.
func (d *ding) askAddressIfNotSet() *ding {
	if d.cliCtx.IsSet("address") {
		d.config.Address = d.cliCtx.String("address")
	} else {
		if err := survey.AskOne(&survey.Input{
			Message: "address to knock:",
		}, &d.config.Address); err != nil {
			return nil
		}
	}

	return d
}

// askPortIfNotSet asks for the port, or sets it if given by flag.
func (d *ding) askPortIfNotSet() *ding {
	if d.cliCtx.IsSet("port") {
		d.config.Port = d.cliCtx.StringSlice("port")
	} else {
		var input string

		if err := survey.AskOne(&survey.Input{
			Message: "port to knock (separated by commas if several):",
		}, &input); err != nil {
			return nil
		}

		d.config.Port = strings.Split(input, ",")
	}

	return d
}

// askTimeoutIfNotSet asks for the timeout, or sets it if given by flag.
func (d *ding) askTimeoutIfNotSet() *ding {
	if d.cliCtx.IsSet("timeout") {
		d.config.Timeout = d.cliCtx.Duration("timeout")
	} else {
		if err := survey.AskOne(&survey.Input{
			Message: "timeout in milliseconds:",
			Default: d.cliCtx.Duration("timeout").String(),
			Help:    "TCP connection timeout, this value should be increased on a slow network",
		}, &d.config.Timeout); err != nil {
			return nil
		}
	}

	return d
}

// askDelayIfNotSet asks for the delay, or sets it if given by flag.
func (d *ding) askDelayIfNotSet() *ding {
	if d.cliCtx.IsSet("delay") {
		d.config.Delay = d.cliCtx.Duration("delay")
	} else {
		if err := survey.AskOne(&survey.Input{
			Message: "delay in milliseconds between knocks:",
			Default: d.cliCtx.Duration("delay").String(),
			Help:    "waiting time between TCP connections, value tied to server configuration",
		}, &d.config.Delay); err != nil {
			return nil
		}
	}

	return d
}

// write creates configuration file and encodes config to it.
func (d *ding) write() error {
	if err := d.createConfigurationFile(); err != nil {
		return err
	}

	if err := toml.NewEncoder(d.configFile).Encode(d.config); err != nil {
		return err
	}

	return d.configFile.Close()
}

// marshal encodes config to a buffer of bytes suitable for post-manipulations.
// It also generates a 128-bit salt, a 256-bit key derived via Argon2id, initializes a XChaCha20-Poly1305 AEAD cipher.
// Then, the buffer is ciphered and written to the configuration file.
func (d *ding) marshal() error {
	if err := d.createConfigurationFile(); err != nil {
		return err
	}

	var dst bytes.Buffer

	if err := toml.NewEncoder(&dst).Encode(d.config); err != nil {
		return err
	}

	salt := make([]byte, 16)

	if _, err := rand.Read(salt); err != nil {
		return nil
	}

	path := d.saltPath()
	if path == "" {
		return errors.New("failed to retrieve salt path")
	}

	encodedSalt := make([]byte, hex.EncodedLen(len(salt)))

	hex.Encode(encodedSalt, salt)

	if err := os.WriteFile(path, encodedSalt, 0644); err != nil {
		return err
	}

	aead, err := chacha20poly1305.NewX(d.deriveKey(salt))
	if err != nil {
		return err
	}

	nonce := make([]byte, aead.NonceSize(), aead.NonceSize()+dst.Len()+aead.Overhead())

	if _, err = rand.Read(nonce); err != nil {
		return err
	}

	if _, err = d.configFile.Write(aead.Seal(nonce, nonce, dst.Bytes(), nil)); err != nil {
		return err
	}

	return d.configFile.Close()
}

func (d *ding) askPassword() *bytes.Buffer {
	var input string

	if err := survey.AskOne(&survey.Password{Message: "password:"}, &input); err != nil {
		panic("prompt error:" + err.Error())
	}

	return bytes.NewBufferString(input)
}

func (d *ding) deriveKey(salt []byte) []byte {
	var password *bytes.Buffer

	for {
		password = d.askPassword()

		// Password entropy must be verified only during the setup phase
		if d.cliCtx.Command.Name != "setup" {
			break
		}

		if !d.cliCtx.Bool("bypass-password-entropy") {
			if err := entropy.Validate(password.String(), 65); err != nil {
				log.Warn().Msg(err.Error())
				continue
			}
		}

		break
	}

	defer password.Reset()

	return argon2.IDKey(password.Bytes(), salt, 1, 64*1024, uint8(runtime.GOMAXPROCS(runtime.NumCPU())), chacha20poly1305.KeySize)
}

// read reads an unencrypted configuration file.
func (d *ding) read() error {
	path := d.configPath()
	if path == "" {
		return errors.New("failed to retrieve configuration file path")
	}

	var err error

	d.configFile, err = os.Open(path)
	if err != nil {
		return err
	}

	defer d.configFile.Close()

	_, err = toml.NewDecoder(d.configFile).Decode(&d.config)

	return err
}

// unmarshal deciphers and decodes the configuration file.
func (d *ding) unmarshal() error {
	path := d.configPath()
	if path == "" {
		return errors.New("failed to retrieve configuration file path")
	}

	conf, err := os.ReadFile(path)
	if err != nil {
		return err
	}

	path = d.saltPath()
	if path == "" {
		return errors.New("failed to retrieve salt path")
	}

	var salt []byte

	salt, err = os.ReadFile(path)
	if err != nil {
		return err
	}

	decodedSalt := make([]byte, hex.DecodedLen(len(salt)))

	if _, err = hex.Decode(decodedSalt, salt); err != nil {
		return err
	}

	var aead cipher.AEAD

	aead, err = chacha20poly1305.NewX(d.deriveKey(decodedSalt))
	if err != nil {
		return err
	}

	if len(conf) < aead.NonceSize() {
		return errors.New("ciphered data size is greater than the nonce size")
	}

	nonce, cipheredData := conf[:aead.NonceSize()], conf[aead.NonceSize():]

	var decipheredData []byte

	decipheredData, err = aead.Open(nil, nonce, cipheredData, nil)
	if err != nil {
		return err
	}

	_, err = toml.NewDecoder(bytes.NewBuffer(decipheredData)).Decode(&d.config)

	return err
}

func validatePortRange(ports []string) error {
	for _, port := range ports {
		p, err := strconv.Atoi(port)
		if err != nil {
			return err
		}

		if p < 1 || p > 65535 {
			return fmt.Errorf("%d: invalid port", p)
		}
	}

	return nil
}