package winpwn

import (
	"bytes"
	"debug/pe"
	"errors"
	"io"
	"math"
)

// Section characteristic flags relevant to memory protection
// (IMAGE_SCN_MEM_*, see winnt.h).
const (
	imageSCNMemExecute = 0x20000000
	imageSCNMemRead    = 0x40000000
	imageSCNMemWrite   = 0x80000000
)

// Section wraps debug/pe.Section with the permission/entropy helpers a pwn
// workflow actually needs (find the RWX section, spot the packed one).
type Section struct {
	*pe.Section

	// live is set when this Section belongs to a live-process-backed
	// PEFile. debug/pe.Section.Data() always reads via PointerToRawData
	// (the section's *file* offset), which is wrong once the image is
	// loaded into memory -- SectionAlignment shifts things around relative
	// to FileAlignment. When live is set, Data() reads VirtualSize bytes
	// at VirtualAddress instead, through the same ReaderAt the rest of a
	// live PEFile uses.
	live io.ReaderAt
}

func (s *Section) IsReadable() bool   { return s.Characteristics&imageSCNMemRead != 0 }
func (s *Section) IsWritable() bool   { return s.Characteristics&imageSCNMemWrite != 0 }
func (s *Section) IsExecutable() bool { return s.Characteristics&imageSCNMemExecute != 0 }
func (s *Section) IsRWX() bool        { return s.IsReadable() && s.IsWritable() && s.IsExecutable() }

// Data returns the section's raw bytes. Overrides (shadows)
// debug/pe.Section.Data: see the live field's doc comment for why a
// live-process-backed section needs a different read path.
func (s *Section) Data() ([]byte, error) {
	if s.live == nil {
		return s.Section.Data()
	}
	buf := make([]byte, s.VirtualSize)
	if _, err := s.live.ReadAt(buf, int64(s.VirtualAddress)); err != nil {
		return nil, err
	}
	return buf, nil
}

// Entropy returns the Shannon entropy (0..8 bits/byte) of the section's raw
// data, the standard quick signal for "this is packed/encrypted" (UPX-style
// sections commonly read >7.2).
func (s *Section) Entropy() (float64, error) {
	data, err := s.Data()
	if err != nil {
		return 0, err
	}
	return ShannonEntropy(data), nil
}

// ShannonEntropy computes the byte-level Shannon entropy of data, in bits
// per byte (0 = uniform/empty, 8 = maximally random).
func ShannonEntropy(data []byte) float64 {
	if len(data) == 0 {
		return 0
	}
	var counts [256]int
	for _, b := range data {
		counts[b]++
	}
	entropy := 0.0
	total := float64(len(data))
	for _, c := range counts {
		if c == 0 {
			continue
		}
		freq := float64(c) / total
		entropy -= freq * math.Log2(freq)
	}
	return entropy
}

// Sections returns every section wrapped with the permission/entropy helpers.
func (p *PEFile) Sections() []*Section {
	out := make([]*Section, len(p.File.Sections))
	for i, sec := range p.File.Sections {
		out[i] = p.wrapSection(sec)
	}
	return out
}

// Section looks up a single section by name (e.g. ".text").
func (p *PEFile) Section(name string) (*Section, error) {
	for _, sec := range p.File.Sections {
		if sec.Name == name {
			return p.wrapSection(sec), nil
		}
	}
	return nil, errors.New("section not found: " + name)
}

func (p *PEFile) wrapSection(sec *pe.Section) *Section {
	s := &Section{Section: sec}
	if p.live {
		s.live = p.r
	}
	return s
}

// LikelyPackedSections returns sections whose entropy exceeds threshold
// (0 selects the common UPX-style default of 7.2 bits/byte), the quick
// "is this binary packed" check pwntools has no direct analogue for since
// ELF packers are rarer in CTF practice than UPX-on-Windows.
func (p *PEFile) LikelyPackedSections(threshold float64) ([]*Section, error) {
	if threshold <= 0 {
		threshold = 7.2
	}
	var hits []*Section
	for _, sec := range p.Sections() {
		if sec.Size == 0 {
			continue
		}
		entropy, err := sec.Entropy()
		if err != nil {
			continue
		}
		if entropy >= threshold {
			hits = append(hits, sec)
		}
	}
	return hits, nil
}

// SearchBytes ищет последовательность байт во всех секциях PE-файла
// Возвращает массив RVA (Relative Virtual Address) всех совпадений
func (p *PEFile) SearchBytes(pattern []byte) ([]uint64, error) {
	var results []uint64

	for _, sec := range p.Sections() {
		// Читаем сырые данные секции (live-aware: see Section.Data)
		data, err := sec.Data()
		if err != nil {
			continue // Если секция пустая (например .bss), пропускаем
		}

		offset := 0
		for {
			// Ищем паттерн в оставшейся части данных
			idx := bytes.Index(data[offset:], pattern)
			if idx == -1 {
				break
			}

			// Вычисляем RVA: виртуальный адрес секции + смещение внутри секции
			rva := sec.VirtualAddress + uint32(offset+idx)
			results = append(results, uint64(rva))

			// Сдвигаем offset, чтобы продолжить поиск после текущего совпадения
			offset += idx + 1
		}
	}

	if len(results) == 0 {
		return nil, errors.New("pattern not found in PE file")
	}

	return results, nil
}