package protocol

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"os"
	"strings"
	"time"

	"github.com/vulncheck-oss/go-exploit/output"
	"golang.org/x/net/proxy"
)

// Connections to the remote target with or without encryption depending on the ssl bool.
func MixedConnect(host string, port int, ssl bool) (net.Conn, bool) {
	if ssl {
		return TLSConnect(host, port)
	}

	return TCPConnect(host, port)
}

// Connects to the remote target with encryption.
func TLSConnect(host string, port int) (net.Conn, bool) {
	conn, ok := TCPConnect(host, port)
	if !ok {
		return nil, false
	}

	return tls.Client(conn, &tls.Config{InsecureSkipVerify: true}), true
}

// Connects to a remote target without encryption.
func TCPConnect(host string, port int) (net.Conn, bool) {
	target := fmt.Sprintf("%s:%d", host, port)

	// do we need to use a proxy?
	envProxy := os.Getenv("ALL_PROXY")
	if strings.HasPrefix(envProxy, "socks5://") {
		dialer, err := proxy.SOCKS5("tcp", envProxy[len("socks5://"):], nil, proxy.Direct)
		if err != nil {
			output.PrintfFrameworkError("SOCKS5 error: %s", err.Error())

			return nil, false
		}

		// timeout long hanging connections
		perHost := proxy.NewPerHost(dialer, proxy.Direct)
		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(GlobalCommTimeout)*time.Second)
		defer cancel()
		conn, err := perHost.DialContext(ctx, "tcp", target)
		if err != nil {
			output.PrintfFrameworkError("Connection failed: %s", err.Error())

			return nil, false
		}

		return conn, true
	}

	// no proxy involved
	conn, err := net.DialTimeout("tcp", target, time.Duration(GlobalCommTimeout)*time.Second)
	if err != nil {
		output.PrintFrameworkError("Connection failed: " + err.Error())

		return nil, false
	}

	return conn, true
}

func TCPWrite(conn net.Conn, data []byte) bool {
	written, err := conn.Write(data)
	if err != nil {
		output.PrintFrameworkError("Server write failed: " + err.Error())

		return false
	}
	if written != len(data) {
		output.PrintFrameworkError("Failed to write all data")

		return false
	}

	return true
}

func TCPReadAmount(conn net.Conn, amount int) ([]byte, bool) {
	reply := make([]byte, amount)
	totalRead := 0

	// keep reading until we hit the desired amount (or an error occurs)
	for totalRead < amount {
		count, err := conn.Read(reply[totalRead:])
		if err != nil {
			output.PrintFrameworkError("Failed to read from the socket: " + err.Error())

			return nil, false
		}
		if count == 0 {
			output.PrintFrameworkError("Connection closed.")

			return nil, false
		}
		totalRead += count
	}

	return reply, true
}

// Read an amount and dont log errors if we fail to read from the socket.
func TCPReadAmountBlind(conn net.Conn, amount int) ([]byte, bool) {
	reply := make([]byte, amount)
	totalRead := 0

	// keep reading until we hit the desired amount (or an error occurs)
	for totalRead < amount {
		count, err := conn.Read(reply[totalRead:])
		if err != nil {
			return nil, false
		}
		if count == 0 {
			return nil, false
		}
		totalRead += count
	}

	return reply, true
}