// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tdx

import (
	"bytes"
	"crypto/sha512"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// NewTDEventLogSpecIdHeader parses a TD event log spec ID header from raw bytes
func NewTDEventLogSpecIdHeader(data []byte) *TDEventLogSpecIdHeader {
	var specIdHeader TDEventLogSpecIdHeader
	specIdHeader.AlgorithmId2DigestSize = make(map[uint16]uint16)

	specIdHeader.TDRegisterIndex = binary.LittleEndian.Uint32(data[:4])
	specIdHeader.EventType = binary.LittleEndian.Uint32(data[4:8])
	specIdHeader.DigestCount = binary.LittleEndian.Uint32(data[8:12])
	copy(specIdHeader.ZeroDigest[:], data[12:32])
	copy(specIdHeader.NoIdea[:], data[32:56])
	specIdHeader.AlgorithmNumber = binary.LittleEndian.Uint32(data[56:60])

	index := 60
	for range specIdHeader.AlgorithmNumber {
		algorithmId := binary.LittleEndian.Uint16(data[index : index+2])
		index = index + 2
		digestSize := binary.LittleEndian.Uint16(data[index : index+2])
		index = index + 2
		specIdHeader.AlgorithmId2DigestSize[algorithmId] = digestSize
	}
	specIdHeader.VendorSize = data[index]
	index = index + int(specIdHeader.VendorSize) + 1
	specIdHeader.RawData = data[:index]

	return &specIdHeader
}

// NewTDEventLogEntry parses a single TD event log entry from raw bytes
func NewTDEventLogEntry(data []byte, tdEventLogSpecIdHeader *TDEventLogSpecIdHeader) *TDEventLogEntry {
	var tdEventLogEntry TDEventLogEntry
	var index = 12

	tdEventLogEntry.TDRegisterIndex = binary.LittleEndian.Uint32(data[:4])
	tdEventLogEntry.EventType = binary.LittleEndian.Uint32(data[4:8])
	tdEventLogEntry.DigestCount = binary.LittleEndian.Uint32(data[8:12])

	for range tdEventLogEntry.DigestCount {
		tdEventLogEntry.AlgorithmId = binary.LittleEndian.Uint16(data[index : index+2])
		index = index + 2
		digestSize := tdEventLogSpecIdHeader.AlgorithmId2DigestSize[tdEventLogEntry.AlgorithmId]
		tdEventLogEntry.Digests = append(tdEventLogEntry.Digests, data[index:index+int(digestSize)])
		index = index + int(digestSize)
	}
	tdEventLogEntry.EventSize = binary.LittleEndian.Uint32(data[index : index+4])
	index = index + 4
	tdEventLogEntry.Event = data[index : index+int(tdEventLogEntry.EventSize)]
	index = index + int(tdEventLogEntry.EventSize)
	tdEventLogEntry.RawData = data[:index]

	return &tdEventLogEntry
}

// NewTDEventLog parses event log data into a TDEventLog structure
func NewTDEventLog(eventLog []byte) (*TDEventLog, error) {
	var tdEventLog TDEventLog
	eventLogData := eventLog

	var index = 0
	for index < len(eventLogData) {
		start := index
		rtmr := binary.LittleEndian.Uint32(eventLogData[index : index+4])
		index = index + 4
		if rtmr == 0xFFFFFFFF {
			break
		}
		eventType := binary.LittleEndian.Uint32(eventLogData[index : index+4])
		index = index + 4

		if eventType == EV_NO_ACTION {
			tdEventLog.Header = NewTDEventLogSpecIdHeader(eventLogData[start:])
			index = start + len(tdEventLog.Header.RawData)
		} else {
			tdEventLogEntry := NewTDEventLogEntry(eventLogData[start:], tdEventLog.Header)
			tdEventLog.Entries = append(tdEventLog.Entries, tdEventLogEntry)
			index = start + len(tdEventLogEntry.RawData)
		}
	}

	return &tdEventLog, nil
}

// Replay replays the event log to calculate RTMR values
// Returns the calculated RTMR values [RTMR0, RTMR1, RTMR2, RTMR3]
func (t *TDEventLog) Replay() (*[TDX_RTMR_COUNT][TDX_RTMR_LENGTH_BY_BYTES]byte, error) {
	var rtmrs [TDX_RTMR_COUNT][TDX_RTMR_LENGTH_BY_BYTES]byte

	sha384 := sha512.New384()
	for _, v := range t.Entries {
		// Verify Digest
		switch v.EventType {
		case EV_EFI_VARIABLE_BOOT:
			variableLog := NewVariableLog(v.Event)
			sha384 := sha512.New384()
			sha384.Write(variableLog.VariableData)
			if !bytes.Equal(sha384.Sum(nil), v.Digests[0]) {
				return nil, errors.New("mismatch digests")
			}
		case EV_EFI_VARIABLE_DRIVER_CONFIG:
			fallthrough
		case EV_SEPARATOR:
			fallthrough
		case EV_EFI_ACTION:
			fallthrough
		case EV_EFI_GPT_EVENT:
			fallthrough
		case EV_EFI_VARIABLE_AUTHORITY:
			sha384 := sha512.New384()
			sha384.Write(v.Event)
			if !bytes.Equal(sha384.Sum(nil), v.Digests[0]) {
				return nil, errors.New("mismatch digests")
			}
		// TODO: Add verification for the following events
		case EV_POST_CODE:
		case EV_PLATFORM_CONFIG_FLAGS:
		case EV_IPL: // GRUB
			if bytes.HasPrefix(v.Event, IPL_KEYWORD_KERNEL_CMDLINE) ||
				bytes.HasPrefix(v.Event, IPL_KEYWORD_MODULE_CMDLINE) ||
				bytes.HasPrefix(v.Event, IPL_KEYWORD_GRUB_CMDLINE) {

				index := bytes.IndexByte(v.Event, ' ')
				sha384 := sha512.New384()
				cmdline := v.Event[index+1:]
				if zero := bytes.IndexByte(cmdline, 0x00); zero >= 0 {
					// Firmware hashes the command string including the first NUL terminator,
					// ignoring any padding that may follow in the event buffer.
					cmdline = cmdline[:zero+1]
				}
				sha384.Write(cmdline)
				if !bytes.Equal(sha384.Sum(nil), v.Digests[0]) {
					return nil, errors.New("mismatch digests")
				}
			} else if bytes.Equal(v.Event, IPL_KEYWORD_MOK_LIST_TRUSTED_BYES) {
				if !bytes.Equal(SHA384_DIGEST_ONE, v.Digests[0]) &&
					!bytes.Equal(SHA384_DIGEST_ZERO, v.Digests[0]) {
					return nil, errors.New("mismatch digests")
				}
			} else if bytes.HasPrefix(v.Event, IPL_KEYWORD_GRUB_LINUXEFI) {
				// The measured file: Kernel or Initrd
				_ = string(v.Event[len(IPL_KEYWORD_GRUB_LINUXEFI)+1:])
			} else if bytes.HasPrefix(v.Event, IPL_KEYWORD_GRUB_KERNEL_CMDLINE) {
				// The grub kernel cmdline and grub.cfg
				_ = string(v.Event[len(IPL_KEYWORD_GRUB_KERNEL_CMDLINE)+1:])
			} else if bytes.Equal(v.Event, IPL_KEYWORD_MOK_LIST) || bytes.Equal(v.Event, IPL_KEYWORD_MOK_LISTX) {
				// TODO: how to verify?
			} else {
				return nil, errors.New("unknown data")
			}
		case EV_EVENT_TAG:
			// AAEL events (e.g. gpu/attestation, container/image-pull) — no special handling needed
		case EV_EFI_BOOT_SERVICES_APPLICATION, EV_EFI_BOOT_SERVICES_DRIVER, EV_EFI_RUNTIME_SERVICES_DRIVER:
		case EV_EFI_PLATFORM_FIRMWARE_BLOB2:
		case EV_EFI_HANDOFF_TABLES2:
			// TODO: Read HANDOFF DATA From address binary.LittleEndian.Uint64(v.Event[34:])
		default:
			// Unknown event types are not fatal — digest accumulation still happens below
		}

		sha384.Reset()
		_, err := sha384.Write(rtmrs[v.TDRegisterIndex-1][:])
		if err != nil {
			return nil, err
		}
		sha384.Write(v.Digests[0])
		copy(rtmrs[v.TDRegisterIndex-1][:], sha384.Sum(nil))
	}

	return &rtmrs, nil
}

// NewVariableLog parses UEFI variable log data
func NewVariableLog(data []byte) *VariableLog {
	var variableLog VariableLog

	copy(variableLog.VariableName[:], data[:16])
	variableLog.UnicodeNameLength = binary.LittleEndian.Uint64(data[16:24])
	variableLog.VariableDataLength = binary.LittleEndian.Uint64(data[24:32])

	variableLog.UnicodeName = make([]uint16, variableLog.UnicodeNameLength)
	for i := 0; i < int(variableLog.UnicodeNameLength); i++ {
		variableLog.UnicodeName[i] = uint16(uint8(data[32+i*2])) | uint16(uint8(data[32+i*2+1]))<<8
	}
	variableLog.VariableData = data[32+2*variableLog.UnicodeNameLength:]
	return &variableLog
}

// ReplayEventLog parses and replays the eventlog to verify RTMR values
// Returns the replayed RTMR values [RTMR0, RTMR1, RTMR2, RTMR3]
func ReplayEventLog(eventLogData []byte) (*[4][48]byte, error) {
	// Parse eventlog
	tdEventLog, err := NewTDEventLog(eventLogData)
	if err != nil {
		return nil, fmt.Errorf("failed to parse TD event log: %w", err)
	}

	// Replay to calculate RTMR values
	rtmrs, err := tdEventLog.Replay()
	if err != nil {
		return nil, fmt.Errorf("failed to replay TD event log: %w", err)
	}

	return rtmrs, nil
}

// ExtractUKIHash finds the last EV_EFI_BOOT_SERVICES_APPLICATION event in RTMR1
// and returns its SHA-384 digest as a hex string.
func ExtractUKIHash(data []byte) (string, error) {
	tdEventLog, err := NewTDEventLog(data)
	if err != nil {
		return "", fmt.Errorf("parse eventlog: %w", err)
	}

	var digest []byte
	for _, e := range tdEventLog.Entries {
		if e.EventType == EV_EFI_BOOT_SERVICES_APPLICATION && e.TDRegisterIndex == TDX_RTMR1_INDEX {
			if len(e.Digests) > 0 {
				digest = e.Digests[0]
			}
		}
	}
	if digest == nil {
		return "", fmt.Errorf("UKI event not found in eventlog")
	}
	return fmt.Sprintf("%x", digest), nil
}

// ExtractKernelDigest finds the td_payload entry in CCEL (PAI form) and
// returns its SHA-384 digest as a hex string. Identification:
//   - Event Type: EV_EFI_PLATFORM_FIRMWARE_BLOB2 (0x8000000a)
//   - Target RTMR: RTMR1
//   - Event payload: BlobDescriptionSize (1B) + BlobDescription ("td_payload"
//     with optional trailing NUL) + BlobBase (8B) + BlobLength (8B).
//
// Returns an error if no matching entry is found.
func ExtractKernelDigest(data []byte) (string, error) {
	tdEventLog, err := NewTDEventLog(data)
	if err != nil {
		return "", fmt.Errorf("parse eventlog: %w", err)
	}

	const wantDesc = "td_payload"
	for _, e := range tdEventLog.Entries {
		if e.EventType != EV_EFI_PLATFORM_FIRMWARE_BLOB2 || e.TDRegisterIndex != TDX_RTMR1_INDEX {
			continue
		}
		if len(e.Event) < 1 {
			continue
		}
		n := int(e.Event[0])
		if 1+n > len(e.Event) {
			continue
		}
		desc := strings.TrimRight(string(e.Event[1:1+n]), "\x00")
		if desc != wantDesc {
			continue
		}
		if len(e.Digests) == 0 {
			continue
		}
		return fmt.Sprintf("%x", e.Digests[0]), nil
	}
	return "", fmt.Errorf("td_payload event not found in eventlog")
}

// ExtractKernelCmdlineDigest finds the td_payload_info entry in CCEL (PAI
// form) and returns its SHA-384 digest as a hex string. Identification:
//   - Event Type: EV_PLATFORM_CONFIG_FLAGS (0xa)
//   - Target RTMR: RTMR1
//   - Event payload prefix: NUL-terminated description "td_payload_info"
//
// Returns an error if no matching entry is found.
func ExtractKernelCmdlineDigest(data []byte) (string, error) {
	tdEventLog, err := NewTDEventLog(data)
	if err != nil {
		return "", fmt.Errorf("parse eventlog: %w", err)
	}

	const wantDesc = "td_payload_info"
	for _, e := range tdEventLog.Entries {
		if e.EventType != EV_PLATFORM_CONFIG_FLAGS || e.TDRegisterIndex != TDX_RTMR1_INDEX {
			continue
		}
		zero := bytes.IndexByte(e.Event, 0)
		if zero < 0 {
			continue
		}
		if string(e.Event[:zero]) != wantDesc {
			continue
		}
		if len(e.Digests) == 0 {
			continue
		}
		return fmt.Sprintf("%x", e.Digests[0]), nil
	}
	return "", fmt.Errorf("td_payload_info event not found in eventlog")
}

// ExtractAAELContainerEvents parses AAEL tagged events (EV_EVENT_TAG with
// tagged_id=AAEL) and returns container image measurements. Two formats are
// recognized:
//
//  1. Legacy / UKI form on RTMR2:
//     event_data = "<domain> container/image-pull {\"name\":\"cvm_container_proxy\",\"digest\":\"sha256:...\"}"
//
//  2. PAI form on RTMR3:
//     event_data = "<domain> kangaroo/pull-image {\"reference\":\"<image-ref>\",\"digest\":\"sha256:...\"}"
//
// Returned events carry whichever of Name / Reference the source provided so
// the caller can match against userdata or perform digest-based matching.
func ExtractAAELContainerEvents(data []byte) ([]AAELContainerEvent, error) {
	tdEventLog, err := NewTDEventLog(data)
	if err != nil {
		return nil, fmt.Errorf("parse eventlog: %w", err)
	}

	var events []AAELContainerEvent
	for _, e := range tdEventLog.Entries {
		if e.EventType != EV_EVENT_TAG {
			continue
		}
		if len(e.Event) < AAEL_TAGGED_EVENT_HEADER_SIZE {
			continue
		}

		// Parse TaggedEvent: tagged_id (4B LE) + tagged_event_size (4B LE) + event_data
		tagID := binary.LittleEndian.Uint32(e.Event[0:4])
		if tagID != AAEL_TAG_ID {
			continue
		}
		tagSize := binary.LittleEndian.Uint32(e.Event[4:AAEL_TAGGED_EVENT_HEADER_SIZE])
		if len(e.Event) < AAEL_TAGGED_EVENT_HEADER_SIZE+int(tagSize) {
			continue
		}
		eventData := string(e.Event[AAEL_TAGGED_EVENT_HEADER_SIZE : AAEL_TAGGED_EVENT_HEADER_SIZE+tagSize])

		// eventData = "{domain} {operation} {content}"
		parts := strings.SplitN(eventData, " ", 3)
		if len(parts) < 3 {
			continue
		}
		operation := parts[1]
		content := parts[2]

		if operation != AAEL_CONTAINER_IMAGE_PULL_OPERATION &&
			operation != AAEL_KANGAROO_PULL_IMAGE_OPERATION {
			continue
		}

		var ctr AAELContainerEvent
		if err := json.Unmarshal([]byte(content), &ctr); err != nil {
			return nil, fmt.Errorf("failed to parse AA container event: %w", err)
		}
		if ctr.Digest == "" {
			continue
		}
		if ctr.Name == "" && ctr.Reference == "" {
			continue
		}
		events = append(events, ctr)
	}

	return events, nil
}
