// Package fortinet is a very basic (and incomplete) implementation of Fortinet FGFM protocol
package fortinet

import (
	"bytes"
	"crypto/tls"
	"encoding/binary"
	"net"

	"github.com/vulncheck-oss/go-exploit/output"
	"github.com/vulncheck-oss/go-exploit/protocol"
)

// Creates and sends a Fortinet FGFM message to a FortiManager.
// The format is closed source, but research by BF, Watchtowr, and Rapid7 have helped uncover the basic message header structure:
// [4 bytes of magic header]
// [4 bytes of total request length]
// [n bytes request body data].
func SendFGFMMessage(conn net.Conn, payload string) bool {
	message := make([]byte, 0)
	// add magic header
	message = append(message, []byte("\x36\xe0\x11\x00")...)
	// build the total length field
	totalLengthField := make([]byte, 4)
	length := len(payload) + 8
	binary.BigEndian.PutUint32(totalLengthField, uint32(length))
	message = append(message, totalLengthField...)
	// add payload
	message = append(message, []byte(payload)...)

	return protocol.TCPWrite(conn, message)
}

// Reads response from a FortiManager.
func ReadFGFMMessage(conn net.Conn) ([]byte, bool) {
	magic, ok := protocol.TCPReadAmount(conn, 4)
	if !ok || !bytes.Equal(magic, []byte("\x36\xe0\x11\x00")) {
		output.PrintFrameworkError("Failed to read server response with expected header")

		return nil, false
	}
	size, ok := protocol.TCPReadAmount(conn, 4)
	if !ok {
		output.PrintFrameworkError("Failed to read server response length")

		return nil, false
	}

	readSize := int(binary.BigEndian.Uint32(size))
	data, ok := protocol.TCPReadAmount(conn, readSize-8)
	if !ok {
		output.PrintFrameworkError("Failed to read server response data")

		return nil, false
	}

	return data, true
}

// Fortimanager requires a connecting Fortigate instance to have a cert.
// SSL is optional here so you have the choice to sign the traffic from the go-exploit framework,
// or so you can send the exploit network traffic through a proxy like socat to sign the traffic for you.
// Benefits to this include being able to generate pcaps of the unencrypted traffic
// between go-exploit and your proxy.
// See CVE-2024-47575 for additional information.
func Connect(host string, port int, ssl bool, cert []byte, key []byte) (net.Conn, bool) {
	if ssl {
		cert, err := tls.X509KeyPair(cert, key)
		if err != nil {
			output.PrintFrameworkError("Failed to load x509 Key Pair")
			output.PrintfFrameworkDebug("Failed to load x509 Key Pair with error: %s", err)

			return nil, false
		}
		cfg := &tls.Config{Certificates: []tls.Certificate{cert}, InsecureSkipVerify: true}

		conn, ok := protocol.TCPConnect(host, port)
		if !ok {
			return nil, false
		}

		return tls.Client(conn, cfg), true
	}

	return protocol.TCPConnect(host, port)
}