package macho

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math/rand"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	. "ios-wrapper/pkg/ios"
	"ios-wrapper/pkg/protomodel"
)

// #include "fixups.h"
import "C"

func rewriteLoadcommandsWithoutCodesignature(mc *MachoContext) {
	if mc.Is64bit() {
		mc.file.Seek(int64(Header_size_64), io.SeekStart)
	} else {
		mc.file.Seek(int64(Header_size), io.SeekStart)
	}
	for _, cmd := range mc.commands {
		if cmd.Cmd() == LC_CODE_SIGNATURE {
			continue
		}
		log.WithFields(log.Fields{
			"cmd":     cmd.Cmd(),
			"cmdsize": cmd.Cmdsize(),
			"name":    cmd.Cmdname(),
		}).Trace("Rewrite Load command")
		mc.file.Write(cmd.Serialize(mc))
	}
}

// Remove Codesign data from Mach-O binary
// There are many caveats to this problem.
// Mostly related to how codesign-allocate works.
// Removing codesign data at the end of __LINKEDIT sement
// is probably easier than removing them at an abitrary location.
//
// Assuming the codesign data is at the end of __LINKEDIT segment.
// The binary is probably signed at the last step, which is common.
//
// CODE_SIGNATURE load commands points to the codesign data offset and size.
// __LINKEDIT section points to data offset and size.
// We have:
//
//	linkedit = (section*) LC_SEGMENT.section[0] // name="__LINKEDIT"
//	codesign = (linkedit_data_command*) LC_CODE_SIGNATURE
//	BinarySize = { f.seek(0, SEEKEND); return f.tell() }
//
//	linkedit->fileoff + linkedit->filesize == codesign->dataOff + codesign->dataSize
//	linkedit->fileoff + linkedit->filesize == BinarySize
//
// To remove the codesign data, we truncate the file to remove the codesign data.
// Then we update the linkedit->filesize to linkedit->filesize - codesign->dataSize
// We also fix the header to not read the LC_CODE_SIGNATURE,
// because codesign is the last step, the last load command is LC_CODE_SIGNATURE.
// Fix header->ncmds -= 1, header->sizeofcmds -= sizeof(linkedit_data_command)
//
// There's one more caveat, as mentioned by `insert_dylib`. Codesign data is aligned
// by 0x10. So it would have some padding before codesign data. Now that we have
// removed the data, but the padding still exists and is not prefered by codesign-allocate.
// To fix this issue, **again** we assume that the previous section is the String table
// (referenced by symtab) and increase the size of the string table to the end of data.
//
// A better approach could be an extended search for section pointing to the previous data.
// But general cases are String table.
//
// Implementation warnings:
// We have two implementation designs.
// - Re-parse the load commands and edit the value at offset
// - Use parsed load commands data, edit the values, and rewrite header + load commands
func (mc *MachoContext) RemoveCodesign() bool {
	if !mc.WriteEnabled() {
		return false
	}

	linkedit := mc.FindSegment("__LINKEDIT")
	if linkedit == nil {
		log.Warn("The binary has no __LINKEDIT Segment")
		return false
	}

	var codesign *LinkEdit
	for _, cmd := range mc.Linkedits() {
		if cmd.Cmd() == LC_CODE_SIGNATURE {
			codesign = cmd
			break
		}
	}
	if codesign == nil {
		log.Warn("The binary is not signed, no LC_CODE_SIGNATURE found")
		return false
	}

	filesize := uint64(len(mc.buf))
	linkedit_end := uint64(linkedit.Fileoff()) + linkedit.Filesize()
	codesign_end := uint64(codesign.Dataoff() + codesign.Datasize())

	// linkedit is the last item in Mach-O binary
	if linkedit_end != filesize {
		log.Panic("Link edit is not the last item")
		return false
	}
	// codesign is the last in linkedit
	if linkedit_end != codesign_end {
		log.Panic("Code sign data is not the last item in link edit segment")
		return false
	}

	linkedit_newsize := linkedit.Filesize() - uint64(codesign.Datasize())

	if (mc.symtab.stroff+mc.symtab.strsize)%0x10 == 0 {
		// codesign requires padding of 0x10
		// if strtab+strsize & 0xff < 0x8, it will pad til 0x10
		// when removed codesign, we truncate the file to offset at 0x8 rather at 0x10
	} else {
		linkedit_newsize -= 8
	}

	mc.file.Truncate(int64(uint64(linkedit.Fileoff()) + linkedit_newsize))
	mc.symtab.strsize = uint32(linkedit.Fileoff()+linkedit_newsize) - mc.symtab.stroff

	// fix linkedit section
	if mc.Is64bit() {
		linkedit.(*Segment64).filesize = linkedit_newsize
	} else {
		linkedit.(*Segment32).filesize = uint32(linkedit_newsize)
	}

	// rewrite load commands
	mc.UpdateHeaderRemoveLcmd(codesign.Cmdsize())
	rewriteLoadcommandsWithoutCodesignature(mc)

	// erase old LC_CODE_SIGNATURE data
	old_codesign_lcmd_offset := func() uint64 {
		loadcmd_size := int64(mc.Header().sizeofcmds)
		if mc.Is64bit() {
			return uint64(loadcmd_size) + Header_size_64
		} else {
			return uint64(loadcmd_size) + Header_size
		}
	}()
	// size of codesign = sizeof linkedit_data_command = 4 * 4
	mc.file.WriteAt(make([]byte, 4*4), int64(old_codesign_lcmd_offset))
	// mc.file.Seek(old_codesign_lcmd_offset, io.SeekStart)
	// mc.file.Write(make([]byte, 4*4))

	return true
}

func (mc *MachoContext) RemoveInitFunctions() bool {
	if mc.WriteEnabled() {
		for _, ptr := range mc.InitFunctions() {
			if mc.Is64bit() {
				mc.file.WriteAt([]byte{0, 0, 0, 0, 0, 0, 0, 0}, int64(ptr.offset))
			} else {
				mc.file.WriteAt([]byte{0, 0, 0, 0}, int64(ptr.offset))
			}
			log.WithFields(log.Fields{
				"offset": fmt.Sprintf("0x%x", ptr.Offset()),
				"value":  fmt.Sprintf("0x%x", ptr.Value()),
			}).Debug("Remove init pointer")
		}
		return true
	}
	return false
}

func (mc *MachoContext) RemoveUnnecessaryInfo() bool {
	for _, command := range mc.commands {
		switch command.(type) {
		case *LinkEdit:
			var le = command.(*LinkEdit)
			if le.Cmdname() != "LC_FUNCTION_STARTS" && le.Cmdname() != "LC_DATA_IN_CODE" {
				continue
			}
			var start int64 = int64(le.dataoff)
			var end int64 = start + int64(le.datasize)
			for i := start; i < end; i++ {
				mc.file.WriteAt([]byte{0}, i)
			}
			continue
		default:
			continue
		}
	}
	return false
}

func (mc *MachoContext) AddLoadCmd(lcmd LoadCommand) {
	var offset uint64
	payload := lcmd.Serialize(mc)
	if uint64(len(payload)) > mc.PaddingSize() {
		log.WithFields(log.Fields{
			"cmd":       lcmd,
			"len(cmd)":  len(payload),
			"available": mc.PaddingSize(),
		}).Panic("Not enough space to add load command")
	}

	if mc.Is64bit() {
		offset = Header_size_64 + uint64(mc.header.sizeofcmds)
	} else {
		offset = Header_size + uint64(mc.header.sizeofcmds)
	}
	log.WithFields(log.Fields{
		"cmd":    lcmd.Cmdname(),
		"size":   len(payload),
		"offset": offset,
	}).Debug("Adding Load Command")
	mc.file.WriteAt(payload, int64(offset))
	mc.UpdateHeaderAddLcmd(uint32(len(payload)))
}

func (mc *MachoContext) AddDylib(dylib string) {
	if mc.DylibExisted(dylib) {
		log.WithFields(log.Fields{
			"dylib": dylib,
		}).Debug("Adding dylib but existed")
		return
	}
	name := []byte(dylib)
	name = append(name, 0)
	dylib_lcmd := Dylib{
		c:                     LoadCmd{cmd: LC_LOAD_DYLIB, cmdsize: 0},
		name:                  name,
		nameoff:               24,
		timestamp:             0,
		current_version:       0,
		compatibility_version: 0,
	}

	mc.AddLoadCmd(&dylib_lcmd)
}

func (mc *MachoContext) AddRPath(rpath string) {
	if mc.RPathExisted(rpath) {
		log.WithFields(log.Fields{
			"rpath": rpath,
		}).Debug("Adding rpath but existed")
		return
	}
	path := []byte(rpath)
	path = append(path, 0)
	rpath_lcmd := RPath{
		c:      LoadCmd{cmd: LC_RPATH, cmdsize: 0},
		offset: 12,
		path:   path,
	}
	mc.AddLoadCmd(&rpath_lcmd)
}

func (mc *MachoContext) UpdateHeaderAddLcmd(size uint32) {
	if mc.WriteEnabled() {
		mc.header.ncmds += 1
		mc.header.sizeofcmds += size
		mc.file.WriteAt(mc.header.Serialize(mc), 0)
	}
}

func (mc *MachoContext) UpdateHeaderRemoveLcmd(size uint32) {
	if mc.WriteEnabled() {
		mc.header.ncmds -= 1
		mc.header.sizeofcmds -= size
		mc.file.WriteAt(mc.header.Serialize(mc), 0)
	}
}

func (mc *MachoContext) RemoveBindSymbols() {
	if !mc.WriteEnabled() {
		return
	}

	rand.Seed(time.Now().UnixNano())

	isModernSymbol := mc.dyldinfo == nil
	isLegacySymbol := !isModernSymbol

	if isModernSymbol {
		mc.removeBindSymbolsModern()
	} else {
		mc.removeBindSymbolsLegacy()
	}
	// Objective-C stub replaces main which can only appears in executable
	if mc.Header().IsExecutable() {
		mc.ReworkForObjc()
	}

	for _, symbol := range mc.CollectBindSymbols() {
		if !symbol.SafeForRemoval() {
			continue
		}

		if isLegacySymbol {
			// for legacy resolve the opcodes can be rewritten as 0x00
			mc.file.WriteAt(make([]byte, 8), int64(symbol.file_address))
		} else {
			// for modern resolve the opcodes must not be rewritten as 0x00
			// because it contains 2 types of opcodes, BIND and REBASE
			// we only fix BIND and leave REBASE intact
			// However, each opcodes has a *next* field to the next opcode
			// So if we want to leave the header intact (contains pointers, size)
			// We should rewrite this as REBASE opcode (so no symbol lookup happens)
			// and it can continue

			// we can write random values, because the loader just do
			// (high8 << 56 | target) - mach_header
			// or something, so no symbol lookup and no error at runtime
			target := rand.Int()
			high8 := rand.Int()
			value := C.MakeRebaseFixupOpcode(C.int(symbol.next), C.ulonglong(target), C.ulonglong(high8))
			v := make([]byte, 8)
			mc.byteorder.PutUint64(v, uint64(value))
			mc.file.WriteAt(v, int64(symbol.file_address))
		}
	}
}

func (mc *MachoContext) removeBindSymbolsModern() {
	// we don't mess up the chain
	// we clear the imports table, and the raw opcodes
	// clearing imports table disables static analysis
	// clearing opcodes forces runtime manual mapping

	// imports item are defined by mc.fixups.imports_format
	// basic case is dyld_chained_import, 4 bytes

	start := mc.fixups.dataoff
	size := mc.fixups.datasize
	fixups := new(Fixups)
	fixups.Deserialize(mc, mc.buf[start:start+size])

	start = mc.fixups.dataoff + fixups.imports_offset
	size = fixups.imports_count * 4
	fmt.Printf("// Erase at=0x%x size=0x%x\n", start, size)
	mc.file.WriteAt(make([]byte, size), int64(start))

	// string reference are at the end of this section
	start = mc.fixups.dataoff + fixups.symbols_offset
	size = mc.fixups.Datasize() - fixups.symbols_offset
	fmt.Printf("// Erase at=0x%x size=0x%x\n", start, size)
	mc.file.WriteAt(make([]byte, size), int64(start))

	fixups.imports_count = 0
	mc.file.WriteAt(fixups.Serialize(mc), int64(mc.fixups.dataoff))
}

func (mc *MachoContext) removeBindSymbolsLegacy() {
	start := mc.dyldinfo.lazy_bind_off
	size := mc.dyldinfo.lazy_bind_size
	// set lazy opcodes to 0x00 == DO_BIND
	// but no symbol state to bind
	mc.file.WriteAt(make([]byte, size), int64(start))
}

func (mc *MachoContext) RewriteImportsTable(keepSymbols []string) {
	allSymbols := mc.CollectBindSymbols()
	fixups, fixupsOffset := mc.Fixups()

	importTable := fixups.ImportsOffset(uint32(fixupsOffset))
	symbolTable := fixups.SymbolsOffset(uint32(fixupsOffset))
	symbolTablePtr := symbolTable

	// in removeBindSymbolsModern, we erase these pointers in file
	// but because we keep a few symbols, we need to rewrite the pointers
	// as well as rebuild the import table and strings table, and bind values

	// some symbols are annoyingly distributed by another library
	// dispite the name asking for X, the dyld loads a Y library
	// LC_DYLD_ID of Y is equal to X and dyld can resolve these symbols
	// because we do not search for library using LC_DYLD_ID,
	// paths are mistaken and will not resolve symbols
	//
	// the most common library that has this behavior is libintl
	// and fixing the resolver takes time, we temporarily ignore this library
	// and so we keep symbols referenced by libintl
	intlSymbols := []string{}
	for _, symbol := range allSymbols {
		if symbol.Dylib() == "/usr/local/opt/gettext/lib/libintl.8.dylib" {
			intlSymbols = append(intlSymbols, symbol.Name())
		}
	}
	keepSymbols = append(keepSymbols, intlSymbols...)

	keepCount := uint32(0)
	for _, symbol := range keepSymbols {
		name := symbol
		lib := ""
		parts := strings.Split(symbol, ",")
		if len(parts) == 2 {
			name = parts[0]
			lib = parts[1]
		}

		symbolInfo := (*ImportSymbol)(nil)
		for _, s := range allSymbols {
			if s.Name() != name {
				continue
			}
			if lib == "" || lib == s.Dylib() {
				symbolInfo = s
				break
			}
		}
		if symbolInfo == nil {
			// symbol not found
			continue
		}

		fmt.Printf("keep symbol %s\n", name)
		fmt.Printf("importTable at 0x%x; stringTable at 0x%x\n", importTable, symbolTablePtr)
		fmt.Printf("bind value at 0x%x\n", symbolInfo.file_address)

		// write string to string table
		mc.file.WriteAt([]byte(name), int64(symbolTablePtr))
		// fix bind value
		rebaseOpcodeBytes := make([]byte, 8)
		mc.file.ReadAt(rebaseOpcodeBytes, int64(symbolInfo.file_address))
		rebaseOpcode := mc.byteorder.Uint64(rebaseOpcodeBytes)
		bindOpcode := C.MakeBindFixupOpcodeFromRebase(C.uint64_t(rebaseOpcode), C.uint(keepCount))

		{
			v := make([]byte, 8)
			mc.byteorder.PutUint64(v, uint64(bindOpcode))
			mc.file.WriteAt(v, int64(symbolInfo.file_address))
		}
		// write import data to import table
		entry := C.MakeImportTableEntry(C.uint(symbolInfo.LibOrdinal()), C.uint(symbolTablePtr-symbolTable))
		{
			v := make([]byte, 4)
			mc.byteorder.PutUint32(v, uint32(entry))
			mc.file.WriteAt(v, int64(importTable))
		}

		keepCount += 1
		importTable += 4
		symbolTablePtr += uint32(len(name)) + 1
	}

	fixups.imports_count = keepCount
	mc.file.WriteAt(fixups.Serialize(mc), int64(mc.fixups.dataoff))
}

func (mc *MachoContext) RemoveSymbolTable() {
	// try to remove symtab and dysymtab
	mc.removeSymtabCommand()
	mc.removeDySymtabCommand()
}

func (mc *MachoContext) removeSymtabCommand() {
	ptr := int64(0)
	if mc.Is64bit() {
		ptr, _ = mc.file.Seek(int64(Header_size_64), io.SeekStart)
	} else {
		ptr, _ = mc.file.Seek(int64(Header_size), io.SeekStart)
	}

	var symtab_fix *Symtab
	for _, cmd := range mc.commands {
		if cmd.Cmd() != LC_SYMTAB {
			ptr += int64(cmd.Cmdsize())
			continue
		}
		var symtab = cmd.(*Symtab)
		symtab_fix = symtab

		// erase strings referenced
		start := int64(symtab_fix.stroff)
		size := symtab_fix.strsize
		fmt.Printf("// Erase at=0x%x size=0x%x\n", start, size)
		mc.file.WriteAt(make([]byte, size), start)

		// erase nlist64 symbol items
		start = int64(symtab_fix.symoff)
		size = symtab_fix.nsyms * 16
		fmt.Printf("// Erase at=0x%x size=0x%x\n", start, size)
		mc.file.WriteAt(make([]byte, size), start)

		symtab_fix.symoff = 0
		symtab_fix.nsyms = 0
		symtab_fix.stroff = 0
		symtab_fix.strsize = 0
		mc.file.Seek(ptr, io.SeekStart)
		mc.file.Write(symtab_fix.Serialize(mc))
		break
	}
	mc.file.Seek(0, io.SeekStart)
}

func (mc *MachoContext) removeDySymtabCommand() {
	ptr := int64(0)
	if mc.Is64bit() {
		ptr, _ = mc.file.Seek(int64(Header_size_64), io.SeekStart)
	} else {
		ptr, _ = mc.file.Seek(int64(Header_size), io.SeekStart)
	}
	for _, cmd := range mc.commands {
		if cmd.Cmd() != LC_DYSYMTAB {
			ptr += int64(cmd.Cmdsize())
			continue
		}
		var dysymtab = cmd.(*DySymtab)
		dysymtab_fix := dysymtab
		dysymtab_fix.indirectsymoff = 0
		dysymtab_fix.nindirectsyms = 0
		mc.file.Seek(ptr, io.SeekStart)
		mc.file.Write(dysymtab_fix.Serialize(mc))
	}
}

func (mc *MachoContext) RemoveExportTrie() {
	var start int64
	var size int
	if mc.dyldinfo != nil {
		// legacy export trie
		start = int64(mc.dyldinfo.export_off)
		size = int(mc.dyldinfo.export_size)
		mc.file.WriteAt(make([]byte, size), start)
	} else if mc.exports != nil {
		// modern export trie
		start = int64(mc.exports.dataoff)
		size = int(mc.exports.datasize)
		mc.file.WriteAt(make([]byte, size), start)
	} else {
		// no export trie (??)
		// should never occur unless this binary is modified
	}
}

func (mc *MachoContext) AddSection(segname string, name string, size int) Section {
	// mc.file.WriteAt(mc.header.Serialize(mc), 0)
	var ret Section
	var buffer bytes.Buffer
	for _, command := range mc.commands {
		switch command.(type) {
		case *Segment64:
			var virtualAddr uint64
			var fileOffset uint64

			var segment = command.(*Segment64)
			if bytes.Compare(bytes.Trim(segment.SegName(), "\x00"), []byte(segname)) != 0 {
				buffer.Write(segment.Serialize(mc))
				continue
			} else {
				virtualAddr = segment.Vmaddr()
				fileOffset = segment.Fileoff()
				for _, section := range segment.Sections() {
					virtualAddr += section.Size()
					// if section.Offset() != 0 {
					fileOffset += section.Size()
					// }
				}

				align := uint64(4)
				alignment := align - (fileOffset % align)
				fileOffset += alignment
				virtualAddr += alignment

				enoughSpace := segment.Fileoff()+segment.Filesize() >= fileOffset+uint64(size)
				if !enoughSpace {
					fmt.Println("Not enough space to store saved info in __DATA segment, need resize (not supported now)")
					panic("Not enough space to store saved info in __DATA segment, need resize (not supported now)")
				}
			}

			var section Section64
			section.sectname = make([]byte, 16)
			copy(section.sectname, name)
			section.segname = make([]byte, 16)
			copy(section.segname, segname)
			section.size = uint64(size)
			section.reloff = 0
			section.nreloc = 0
			section.flags = 0
			section.align = 3
			section.reserved1 = 0
			section.reserved2 = 0
			section.reserved3 = 0

			// addr will increment from the last section
			// offset will increment from the last section + size
			// be careful of Virtual section (bss)
			section.addr = virtualAddr
			section.offset = uint32(fileOffset)
			segment.nsects += 1
			buffer.Write(segment.Serialize(mc))
			buffer.Write(section.Serialize(mc))

			fmt.Printf("Add a new section with addr=0x%x, fileoffset=0x%x\n", section.addr, section.offset)

			ret = &section
			continue

		default:
			buffer.Write(command.Serialize(mc))
			continue
		}
	}

	mc.header.sizeofcmds = uint32(buffer.Len())
	header := mc.header.Serialize(mc)
	mc.file.WriteAt(header, 0)
	mc.file.WriteAt(buffer.Bytes(), int64(len(header)))
	return ret
}

func (mc *MachoContext) WriteInfoToData(info *protomodel.MachoInfo) {
	encode := func() []byte {
		buffer := new(bytes.Buffer)
		for _, table := range info.Symbols.Tables {
			binary.Write(buffer, mc.byteorder, table.LibIndex)
			binary.Write(buffer, mc.byteorder, table.Nsymbols)
			for _, symbol := range table.Symbols {
				binary.Write(buffer, mc.byteorder, (symbol.SymbolIndex<<8)|symbol.SegmentIndex)
				binary.Write(buffer, mc.byteorder, symbol.Offset)
			}
		}
		instructions := buffer.Bytes()

		buffer = new(bytes.Buffer)
		for _, lib := range info.Symbols.Libs {
			buffer.WriteString(lib)
			buffer.WriteByte(0)
		}
		liblist := buffer.Bytes()

		buffer = new(bytes.Buffer)
		for _, symbol := range info.Symbols.Symbols {
			buffer.WriteString(symbol)
			buffer.WriteByte(0)
		}
		symbollist := buffer.Bytes()

		buffer = new(bytes.Buffer)
		// offset to liblist
		binary.Write(buffer, mc.byteorder, uint32(len(instructions)))
		// offset to symbollist
		binary.Write(buffer, mc.byteorder, uint32(len(instructions)+len(liblist)))
		buffer.Write(instructions)
		buffer.Write(liblist)
		buffer.Write(symbollist)

		return buffer.Bytes()
	}
	encoded := encode()
	// encoded := []byte{0x11,0x22,0x33, 0x44}
	section := mc.AddSection("__DATA", "selfbind", len(encoded))
	if mc.Is64bit() {
		var s *Section64 = section.(*Section64)
		mc.file.WriteAt(encoded, int64(s.Offset()))
	} else {
		var s *Section32 = section.(*Section32)
		mc.file.WriteAt(encoded, int64(s.Offset()))
	}
}