// Package ikev2
// A (very) basic framework for an ikev2 protocol.
// The intent with adding this package is not to add an entire, fully functional protocol but rather providing a convenient interface for generating the data structures for ike-related exploits.
// The goal is to add to this as we go, as new exploits demand it.
//
// Using this will probably require you to make your own SAInit function to generate the message data and then sending it over the ikeClient connection.
// A basic example of something like this follows:
//
//	func saInit(ikeClient *ikev2.IkeClient, diffieHellmanGroup int) ([]byte, bool) {
//
//	   ikeClient.IkeCrypto = ikev2.IkeCrypto{}
//	   ikeClient.IkeCrypto.Init()
//	   ok := ikeClient.IkeCrypto.GenerateDHKey(diffieHellmanGroup)
//	   if !ok {
//	   	return []byte{}, false
//	   }
//
//	   defaultTransforms := []ikev2.IkeTransform{
//	   	{NextPayload: ikev2.PayloadType["TRANSFORM"], TransformType: ikev2.TransformType["ENCRYPTION_ALGORITHM"], TransformID: ikev2.EncryptionAlgorithm["ENCR_AES_CBC"], TransformAttributes: 0x800e0100},
//	   }
//
//	   message := make([]byte, 0)
//	   header := ikev2.IkePackHeader(ikeClient, ikev2.PayloadType["SECURITY_ASSOCIATION"], 0x20, ikev2.ExchangeType["IKE_SA_INIT"], 0x08, 0x0)
//	   message = append(message, ikev2.IkePackSecurityAssociation(ikev2.PayloadType["KEY_EXCHANGE"], ikev2.IkePackProposal(ikev2.PayloadType["NONE"], 1, 1, defaultTransforms, ""))...)
//	   message = append(message, ikev2.IkePackKeyExchange(ikev2.PayloadType["NONCE"], ikev2.DiffieHellmanGroup["DH_GROUP_2048_BIT_MODP"], ikeClient.IkeCrypto.ClientPubKey.Bytes())...)
//	   message = append(message, ikev2.IkePackNonce(ikev2.PayloadType["NOTIFY"], ikeClient.IkeCrypto.InitNonce)...)
//
//	   tempBytes, _ := hex.DecodeString("deadbeef")
//	   message = append(message, ikev2.IkePackNotify(ikev2.PayloadType["NOTIFY"], ikev2.NotifyType["NAT_DETECTION_DESTINATION_IP"], tempBytes, 1, 0)...)
//
//	   tempBytes, _ = hex.DecodeString("0021030203030006")
//	   message = append(message, ikev2.IkePackNotify(ikev2.PayloadType["NONE"], ikev2.NotifyType["SIGNATURE_HASH_ALGORITHMS"], tempBytes, 0, 0)...)
//
//	   // combining the message
//	   payload := make([]byte, 0)
//	   payload = append(header, []byte(transform.PackBigInt32(len(message)+len(header)+4))...)
//	   payload = append(payload, message...)
//	   return payload, true
//	}
//
//	func main() {
//	   ikeClient := ikev2.IkeClient{}
//	   saInitPayload, _ := saInit(&ikeClient, ikev2.DiffieHellmanGroup["DH_GROUP_2048_BIT_MODP"])
//	   if !ok {
//	   	output.PrintError("error encountered while generating SAInit payload")
//	   	return false
//	   }
//
//	   if !ikeClient.Connect(conf){
//	   	return false
//	   }
//
//	   _, err := ikeClient.Conn.Write([]byte(saInitPayload))
//	   if err != nil {
//	   	output.PrintfDebug("SAInit Send failed: %s", err.Error())
//	   	return false
//	   }
//	   output.PrintDebug("Sent SAInit message")
//	}
package ikev2

import (
	"crypto/rand"
	"encoding/binary"
	"math/big"
	mrand "math/rand"
	"net"
	"strconv"

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

func Uint64ToString(num uint64) []byte {
	bytes := make([]byte, 8)
	binary.BigEndian.PutUint64(bytes, num)

	return bytes
}

func generateNonce(length int) ([]byte, bool) {
	nonce := make([]byte, length)
	_, err := rand.Read(nonce)
	if err != nil {
		output.PrintfFrameworkError("Error while generating nonce: %s", err.Error())

		return []byte{}, false
	}

	return nonce, true
}

// Switch case is here in case anyone wants to add additional DH support.
func (ikeCrypto *IkeCrypto) GenerateDHKey(diffieHellmanGroup int) bool {
	switch diffieHellmanGroup {
	case DiffieHellmanGroup["DH_GROUP_2048_BIT_MODP"]:
		ikeCrypto.Prime = new(big.Int)
		ikeCrypto.Prime.SetString("FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F83655D23DCA3AD961C62F356208552BB9ED529077096966D670C354E4ABC9804F1746C08CA18217C32905E462E36CE3BE39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF6955817183995497CEA956AE515D2261898FA051015728E5A8AACAA68FFFFFFFFFFFFFFFF", 16) // https://www.ietf.org/rfc/rfc3526.txt
		g := big.NewInt(2)
		ikeCrypto.ClientPrivKey, _ = rand.Int(rand.Reader, ikeCrypto.Prime)
		ikeCrypto.ClientPubKey = new(big.Int).Exp(g, ikeCrypto.ClientPrivKey, ikeCrypto.Prime)

		return true

	default:
		output.PrintFrameworkError("Provided DH group is currently unsupported.")

		return false
	}
}

func (ikeCrypto *IkeCrypto) Init() bool {
	var ok bool
	ikeCrypto.InitNonce, ok = generateNonce(32)
	if !ok {
		return false
	}
	ikeCrypto.RespNonce = []byte{}
	ikeCrypto.InitSPI = mrand.Uint64()
	ikeCrypto.RespSPI = 0x0000000000000000

	return true
}

// An initial reset/connect helper function.
func (ikeClient *IkeClient) Connect(conf *config.Config) bool {
	var err error
	if ikeClient.Conn != nil {
		err = ikeClient.Conn.Close()
		if err != nil {
			// we probably do not want to exit on this, we're establishing a new connection anyways
			output.PrintfFrameworkWarn("Failed to close socket: %s", err.Error())
		}
	}

	connectionString := net.JoinHostPort(conf.Rhost, strconv.Itoa(conf.Rport))
	ikeClient.Conn, err = net.Dial("udp", connectionString)
	if err != nil {
		output.PrintfFrameworkError("Dial failed: %s", err.Error())

		return false
	}

	output.PrintfFrameworkDebug("Successfully established UDP connection to %s", connectionString)

	return true
}