package regparser

/* This file contains helper methods attached to the autogenerated structs. */

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"time"

	"github.com/davecgh/go-spew/spew"
)

func filetimeToUnixtime(ft uint64) uint64 {
	return (ft - 11644473600000*10000) / 10000000
}

// A FileTime object is a timestamp in windows filetime format.
type FileTime struct {
	time.Time
}

func (self *FileTime) GoString() string {
	return fmt.Sprintf("%v", self)
}

func (self *FileTime) DebugString() string {
	return fmt.Sprintf("%v", self)
}

func (self *RegistryProfile) FileTime(reader io.ReaderAt, offset int64) *FileTime {
	filetime := ParseUint64(reader, offset)
	return &FileTime{time.Unix(int64(filetimeToUnixtime(filetime)), 0)}
}

// UTF16 null terminated string.
type UnicodeString struct {
	Value string
}

func (self *UnicodeString) DebugString() string {
	return fmt.Sprintf("%v", self)
}

func (self *RegistryProfile) UnicodeString(reader io.ReaderAt, offset int64) *UnicodeString {
	result := UnicodeString{""}
	buff := make([]byte, 256)
	_, err := reader.ReadAt(buff, offset)
	if err != nil {
		return &result
	}

	utf8 := UTF16BytesToUTF8(buff, binary.LittleEndian)
	var i int
	var c rune
	for i, c = range utf8 {
		if c == 0 {
			break
		}
	}
	result.Value = utf8[:i]

	return &result
}

func (self *UnicodeString) GoString() string {
	return self.Value
}

// HBASE_BLOCK is the file header block at the start of the registry file.
func (self *HBASE_BLOCK) HiveBin() *HBIN {
	return self.Profile.HBIN(self.Reader, self.Offset+int64(self.Size()))
}

// All data in the registry file is contained in cells. The HCELL
// struct is the main container for everything. We add many
// convenience methods on this structure to be able to extract the
// various things contained inside the cell.
func (self *HCELL) NextCell() *HCELL {
	next := self.Next()
	if self.Allocated() {
		next = -self.Next()
	}
	return self.Profile.HCELL(self.Reader, self.Offset+int64(next))
}

// Cells may be allocated or not.
func (self *HCELL) Allocated() bool {
	return self.Next() > 0x80000000
}

// This method returns the actual size of the cell's data payload.
func (self *HCELL) DataSize() uint32 {
	// https://github.com/msuhanov/regf: Size of a current cell in
	// bytes, including this field (aligned to 8 bytes): the size
	// is positive if a cell is unallocated or negative if a cell
	// is allocated
	cell_size := self.Next()
	cell_size = -cell_size - 4
	cell_size -= cell_size % 8
	return cell_size
}

// The offset of the cells payload.
func (self *HCELL) Payload() int64 {
	return self.Offset + self.Profile.Off_HCELL_Data
}

// If the HCELL contains a CM_KEY_NODE (nk node) then this method
// returns it. Otherwise it returns nil.
func (self *HCELL) KeyNode() *CM_KEY_NODE {
	if self.Signature() == 0x6b6e /* nk */ {
		return self.Profile.CM_KEY_NODE(self.Reader, self.Offset+4)
	}

	return nil
}

// If the HCELL contains a CM_KEY_INDEX (ri or li node) then this method
// returns it. Otherwise it returns nil.
func (self *HCELL) KeyIndex() *CM_KEY_INDEX {
	switch self.Signature() {
	case 0x6972 /* ri */, 0x696c /* li */ :
		return self.Profile.CM_KEY_INDEX(self.Reader, self.Offset+4)
	}

	return nil
}

// If the HCELL contains a CM_KEY_INDEX_FAST (lf or lh node) then this
// method returns it. Otherwise it returns nil.
func (self *HCELL) KeyIndexFast() *CM_KEY_INDEX_FAST {
	switch self.Signature() {
	case 0x666c /* lf */, 0x686c /* lh */ :
		return self.Profile.CM_KEY_INDEX_FAST(self.Reader, self.Offset+4)
	}
	return nil
}

// If the HCELL contains a CM_KEY_VALUE (vk node) then this method
// returns it. Otherwise it returns nil.
func (self *HCELL) KeyValue() *CM_KEY_VALUE {
	switch self.Signature() {
	case 0x6b76 /* vk */ :
		return self.Profile.CM_KEY_VALUE(self.Reader, self.Offset+4)
	}

	return nil
}

// This is a convenience method for enumerating the subkeys of a
// CM_KEY_NODE. Each _CM_KEY_NODE can point to a number of different
// types of index nodes. This method deals with the different types of
// indexes and just returns a list of subkeys regardless of the type
// of indexes.
func (self *CM_KEY_NODE) Subkeys() []*CM_KEY_NODE {
	DebugPrint("_CM_KEY_NODE.Subkeys %x @ %x", self.Signature(), self.Offset)
	child := self.Profile.HCELL(
		self.Reader, 4096+int64(self.SubKeyLists()[0]))

	if child.Signature() == 0 {
		return nil
	}

	// It can point to a fast index of its subkeys.
	if fl := child.KeyIndexFast(); fl != nil {
		return fl.Subkeys()

		// Or one of the regular indexes.
	} else if fi := child.KeyIndex(); fi != nil {
		return fi.Subkeys()

	} else {
		DebugPrint("Unknown KeyNode child %x @ %x",
			child.Signature(), child.Offset)
	}

	return nil
}

// A convenience method for extracting the Values contained under a
// key.
func (self *CM_KEY_NODE) Values() []*CM_KEY_VALUE {
	result := []*CM_KEY_VALUE{}
	count := self.ValueList().Count()

	value_list_cell := self.Profile.HCELL(
		self.Reader,
		4096+int64(self.ValueList().List()))

	value_offsets := ParseSafeArray_uint32(
		self.Reader,
		value_list_cell.Payload(), int(count))

	for _, offset := range value_offsets {
		node := self.Profile.HCELL(self.Reader, 4096+int64(offset))

		if value_node := node.KeyValue(); value_node != nil {
			result = append(result, value_node)
		}
	}

	return result
}

// The name of the a key. This does not include the full path through
// its parents.
func (self *CM_KEY_NODE) Name() string {
	// NameLength is potentially corrupt, but since it is a uint16, this can
	// at worst lead to a 64 KB allocation, which is acceptable.
	buff := make([]byte, self.NameLength())
	n, err := self.Reader.ReadAt(
		buff, self.Profile.Off_CM_KEY_NODE__Name+self.Offset)
	if err == nil {
		return string(buff[:n])
	}
	return ""
}

// Extract all subkeys stored in the fast index.
func (self *CM_KEY_INDEX_FAST) Subkeys() []*CM_KEY_NODE {
	DebugPrint("_CM_KEY_INDEX_FAST.Subkeys %x @ %x", self.Signature(), self.Offset)
	result := []*CM_KEY_NODE{}

	for _, element := range ParseSafeArray_CM_KEY_INDEX_FAST_ELEMENT(
		self.Profile,
		self.Reader,
		self.Offset+self.Profile.Off_CM_KEY_INDEX_FAST_List,
		int(self.Count())) {
		child := self.Profile.HCELL(self.Reader, 4096+int64(
			element.NodeOffset()))

		if child_key := child.KeyNode(); child_key != nil {
			result = append(result, child_key)

		} else {
			DebugPrint("%x: Unknown %x @ %x", self.Signature(),
				child.Signature(), child.Offset)
		}
	}

	return result
}

// Extract subkeys from the index.
func (self *CM_KEY_INDEX) Subkeys() []*CM_KEY_NODE {
	DebugPrint("_CM_KEY_INDEX.Subkeys %x @ %x", self.Signature(), self.Offset)

	result := []*CM_KEY_NODE{}

	for _, offset := range ParseSafeArray_uint32(
		self.Reader,
		self.Offset+self.Profile.Off_CM_KEY_INDEX_List,
		int(self.Count())) {
		child := self.Profile.HCELL(self.Reader, 4096+int64(offset))
		if child_key := child.KeyNode(); child_key != nil {
			result = append(result, child_key)

		} else if child_idx := child.KeyIndex(); child_idx != nil {
			result = append(result, child_idx.Subkeys()...)

		} else if child_idx := child.KeyIndexFast(); child_idx != nil {
			result = append(result, child_idx.Subkeys()...)

		} else {
			DebugPrint("%x: Unknown %x @ %x", self.Signature(),
				child.Signature(), child.Offset)
		}
	}

	return result
}

// The name of this value (empty string means default value).
func (self *CM_KEY_VALUE) ValueName() string {
	// NameLength is potentially corrupt, but since it is a uint16, this can
	// at worst lead to a 64 KB allocation, which is acceptable.
	buff := make([]byte, self.NameLength())
	n, err := self.Reader.ReadAt(
		buff, self.Profile.Off_CM_KEY_VALUE_Name+self.Offset)
	if err == nil {
		return string(buff[:n])
	}
	return ""
}

// A Registry Value may represent a number of different data types
// depending on its Type field. This struct contains the various Go
// types that are represented. Many of the registry types are
// converted to the most closely matching Go types. The original
// binary data is also attached in the Data field.
type ValueData struct {
	// REG_SZ etc.
	Type uint32

	// Filled in for REG_SZ etc.
	String string

	// Filled in for REG_MULTI_SZ
	MultiSz []string

	// Filled in for integer types
	Uint64 uint64

	// The original encoded data. For BINARY_SZ this is the only
	// field filled.
	Data []byte

	// If an error occurs during parsing this will contain the
	// error object.
	Error error
}

// Convert the registry type to a string.
func (self *CM_KEY_VALUE) TypeString() string {
	return RegTypeToString(self.Type())
}

func (self *CM_KEY_VALUE) DataSize() int64 {
	data_size := self.DataLength()
	if data_size&0x80000000 > 0 {
		// Data is so short it can be stored in the Data()
		// field itself. In this case Data() is not a pointer
		// but the actual data value.
		data_size ^= 0x80000000
	}

	return int64(data_size)
}

// Parse out the data from the value into a Go ValueData type.
func (self *CM_KEY_VALUE) ValueData() *ValueData {
	result := &ValueData{Type: self.Type()}

	data_size := self.DataLength()

	if data_size&0x80000000 > 0 {
		// Data is so short it can be stored in the Data()
		// field itself. In this case Data() is not a pointer
		// but the actual data value.
		data_size ^= 0x80000000
		result.Data = ParseSafeArray_byte(
			self.Reader,
			self.Offset+self.Profile.Off_CM_KEY_VALUE_Data, 4)
	} else {
		// Handle BIG_DATA values. The Data field is a pointer
		// to the cell that contains the CM_BIG_DATA object.
		cell := self.Profile.HCELL(self.Reader,
			0x1000+int64(self.Data()))

		if cell.Signature() == 0x6264 /* db */ {
			big_data := self.Profile.CM_BIG_DATA(
				self.Reader, cell.Payload())

			// Big data's List member is a pointer to the
			// CELL that contains the segment list.
			list_cell := self.Profile.HCELL(
				self.Reader, 0x1000+int64(big_data.List()))

			// Extract the segment offset list.
			segment_list := ParseSafeArray_uint32(
				self.Reader,
				list_cell.Payload(),
				int(big_data.Count()))

			for _, offset := range segment_list {
				// Each segment offset is a pointer to
				// the cell that contains the value's
				// data.
				segment_cell := self.Profile.HCELL(self.Reader,
					0x1000+int64(offset))

				// Unallocated cells are probably
				// deleted values.
				if !segment_cell.Allocated() {
					continue
				}

				// There isnt an actual header for
				// these segments, the data is just
				// stored in the cell's payload.
				segment_cell_size := segment_cell.DataSize()
				if segment_cell_size > data_size {
					segment_cell_size = data_size
				}

				result.Data = append(result.Data,
					ParseSafeArray_byte(
						self.Reader,
						segment_cell.Payload(),
						int(segment_cell_size))...)

				// Keep reading more segments until we
				// have all the required length.
				data_size -= segment_cell_size
			}
		} else {
			// Data is stored directly in a single cell.
			result.Data = ParseSafeArray_byte(
				self.Reader,
				cell.Payload(), int(data_size))
		}
	}

	switch result.Type {
	case REG_SZ, REG_EXPAND_SZ:
		result.String = UTF16BytesToUTF8(result.Data, binary.LittleEndian)

	case REG_MULTI_SZ:
		// Make sure the data is valid - len must be even.
		if len(result.Data)%2 > 0 {
			result.Data = append(result.Data, 0)
		}
		var tmp []byte
		for i := 0; i < len(result.Data); i += 2 {
			// NULL character represents the end of one string - flush
			// it and reset for the next string.
			if result.Data[i] == 0 && result.Data[i+1] == 0 {
				if len(tmp) > 0 {
					result.MultiSz = append(result.MultiSz,
						UTF16BytesToUTF8(tmp, binary.LittleEndian))
				}
				tmp = nil
			} else {
				tmp = append(tmp, result.Data[i], result.Data[i+1])
			}
		}

	case REG_BINARY:
		// Leave the raw data as is in the ValueData struct.

	case REG_DWORD:
		if data_size != 4 {
			result.Error = errors.New("Data is of incorrect size")
		} else {
			result.Uint64 = uint64(binary.LittleEndian.Uint32(result.Data))
		}

	case REG_DWORD_BIG_ENDIAN:
		if data_size != 4 {
			result.Error = errors.New("Data is of incorrect size")
		} else {
			result.Uint64 = uint64(binary.BigEndian.Uint32(result.Data))
		}

	case REG_QWORD:
		if data_size != 8 {
			result.Error = errors.New("Data is of incorrect size")
		} else {
			result.Uint64 = binary.LittleEndian.Uint64(result.Data)
		}

	default:
		// We cant handle this data directly - just pass it to
		// the consumer of the ValueData struct to deal with
		// parsing it.
	}

	return result
}

func (self *ValueData) GoString() string {
	switch self.Type {
	case REG_SZ, REG_EXPAND_SZ:
		return self.String

	case REG_DWORD:
		return fmt.Sprintf("%x", self.Uint64)

	case REG_BINARY:
		data := self.Data
		if len(data) > 1024 {
			data = data[:1024]
		}
		return spew.Sdump(data)
	}

	return spew.Sdump(self)
}

func ParseSafeArray_uint32(reader io.ReaderAt, offset int64, count int) []uint32 {
	result := []uint32{}
	data := make([]byte, 4)
	for i := 0; i < count; i++ {
		_, err := reader.ReadAt(data, offset)
		if err == io.EOF {
			break
		} else if err != nil {
			continue
		}
		result = append(result, binary.LittleEndian.Uint32(data))
		offset += 4
	}
	return result
}

func ParseSafeArray_byte(reader io.ReaderAt, offset int64, count int) []byte {
	result := []byte{}
	var data [1]byte
	for i := 0; i < count; i++ {
		_, err := reader.ReadAt(data[:], offset)
		if err == io.EOF {
			break
		} else if err != nil {
			continue
		}
		result = append(result, data[0])
		offset += 1
	}
	return result
}

func ParseSafeArray_CM_KEY_INDEX_FAST_ELEMENT(profile *RegistryProfile, reader io.ReaderAt, offset int64, count int) []*CM_KEY_INDEX_FAST_ELEMENT {
	result := []*CM_KEY_INDEX_FAST_ELEMENT{}
	var probe [1]byte
	for i := 0; i < count; i++ {
		// Do a single byte read to probe whether we are still at a valid position
		_, err := reader.ReadAt(probe[:], offset)
		if err == io.EOF {
			break
		} else if err != nil {
			continue
		}
		value := profile.CM_KEY_INDEX_FAST_ELEMENT(reader, offset)
		result = append(result, value)
		offset += int64(value.Size())
	}
	return result
}
