/*
Copyright © 2018-2026 blacktop

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/
package dyld

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"

	"github.com/apex/log"
	"github.com/blacktop/go-macho"
	"github.com/blacktop/go-macho/pkg/fixupchains"
	"github.com/blacktop/ipsw/internal/profile"
	"github.com/blacktop/ipsw/pkg/dyld"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/vbauerster/mpb/v8"
	"github.com/vbauerster/mpb/v8/decor"
)

var dyldExtractProfileFlags profile.ProfilingFlags

// Keep extract profiling off until these internal perf knobs are ready for the public CLI.
const enableDyldExtractProfiling = false

func rebaseMachO(dsc *dyld.File, machoPath string) error {
	f, err := os.OpenFile(machoPath, os.O_RDWR, 0755)
	if err != nil {
		return fmt.Errorf("failed to open exported MachO %s: %v", machoPath, err)
	}
	defer f.Close()

	mm, err := macho.NewFile(f)
	if err != nil {
		return err
	}

	for _, seg := range mm.Segments() {
		uuid, mapping, err := dsc.GetMappingForVMAddress(seg.Addr)
		if err != nil {
			return err
		}

		if mapping.SlideInfoOffset == 0 {
			continue
		}

		startAddr := seg.Addr - mapping.Address
		endAddr := ((seg.Addr + seg.Memsz) - mapping.Address) + uint64(dsc.SlideInfo.GetPageSize())

		start := startAddr / uint64(dsc.SlideInfo.GetPageSize())
		end := endAddr / uint64(dsc.SlideInfo.GetPageSize())

		rebases, err := dsc.GetRebaseInfoForPages(uuid, mapping, start, end)
		if err != nil {
			return err
		}

		for _, rebase := range rebases {
			off, err := mm.GetOffset(rebase.CacheVMAddress)
			if err != nil {
				continue
			}
			if _, err := f.Seek(int64(off), io.SeekStart); err != nil {
				return fmt.Errorf("failed to seek in exported file to offset %#x from the start: %v", off, err)
			}
			if err := binary.Write(f, dsc.ByteOrder, rebase.Target); err != nil {
				return fmt.Errorf("failed to write rebase address %#x: %v", rebase.Target, err)
			}
		}
	}

	return nil
}

func init() {
	DyldCmd.AddCommand(dyldExtractCmd)
	dyldExtractCmd.Flags().BoolP("all", "a", false, "Extract all dylibs")
	dyldExtractCmd.Flags().Bool("force", false, "Overwrite existing extracted dylib(s)")
	dyldExtractCmd.Flags().Bool("slide", false, "Apply slide info to extracted dylib(s)")
	dyldExtractCmd.Flags().Bool("objc", false, "Add ObjC metadata to extracted dylib(s) symtab (implies --slide)")
	dyldExtractCmd.Flags().Bool("stubs", false, "Add stub islands to extracted dylib(s) symtab")
	dyldExtractCmd.Flags().StringP("cache", "c", "", "Path to .a2s addr to sym cache file (speeds up analysis)")
	dyldExtractCmd.Flags().StringP("output", "o", "", "Directory to extract the dylib(s)")
	dyldExtractCmd.MarkFlagDirname("output")
	if enableDyldExtractProfiling {
		profile.AddFlags(dyldExtractCmd, &dyldExtractProfileFlags)
	}
	viper.BindPFlag("dyld.extract.all", dyldExtractCmd.Flags().Lookup("all"))
	viper.BindPFlag("dyld.extract.force", dyldExtractCmd.Flags().Lookup("force"))
	viper.BindPFlag("dyld.extract.slide", dyldExtractCmd.Flags().Lookup("slide"))
	viper.BindPFlag("dyld.extract.objc", dyldExtractCmd.Flags().Lookup("objc"))
	viper.BindPFlag("dyld.extract.stubs", dyldExtractCmd.Flags().Lookup("stubs"))
	viper.BindPFlag("dyld.extract.cache", dyldExtractCmd.Flags().Lookup("cache"))
	viper.BindPFlag("dyld.extract.output", dyldExtractCmd.Flags().Lookup("output"))
}

// dyldExtractCmd represents the extractDyld command
var dyldExtractCmd = &cobra.Command{
	Use:     "extract <DSC> <DYLIB>",
	Aliases: []string{"e"},
	Short:   "Extract dylib(s) for reverse engineering",
	Args:    cobra.MinimumNArgs(1),
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		if len(args) == 1 {
			return getImages(args[0]), cobra.ShellCompDirectiveDefault
		}
		return getDSCs(toComplete), cobra.ShellCompDirectiveDefault
	},
	SilenceErrors: true,
	RunE: func(cmd *cobra.Command, args []string) error {

		var bar *mpb.Bar
		var p *mpb.Progress
		var images []*dyld.CacheImage

		// flags
		dumpALL := viper.GetBool("dyld.extract.all")
		forceExtract := viper.GetBool("dyld.extract.force")
		slide := viper.GetBool("dyld.extract.slide")
		addObjc := viper.GetBool("dyld.extract.objc")
		addStubs := viper.GetBool("dyld.extract.stubs")
		output := viper.GetString("dyld.extract.output")
		cacheFile := viper.GetString("dyld.extract.cache")
		// --objc writes resolved ObjC symbol addresses, so the dylib's own
		// pointers must be rebased too; otherwise the extracted file holds raw
		// cache slide-info pointers and cannot be re-parsed (e.g. by class-dump).
		if addObjc {
			slide = true
		}
		if dumpALL && len(args) > 1 {
			return fmt.Errorf("cannot specify DYLIB(s) when using --all")
		} else if !dumpALL && len(args) < 2 {
			return fmt.Errorf("must specify at least one DYLIB to extract")
		}

		if enableDyldExtractProfiling {
			prof := profile.New(dyldExtractProfileFlags.ToConfig())
			if err := prof.Start(); err != nil {
				return fmt.Errorf("failed to start profiling: %v", err)
			}
			defer func() {
				if err := prof.Stop(); err != nil {
					log.Errorf("failed to stop profiling: %v", err)
				}
				if dyldExtractProfileFlags.IsEnabled() {
					prof.PrintStats()
				}
			}()
		}

		dscPath := filepath.Clean(args[0])

		folder := filepath.Dir(dscPath) // default to folder of shared cache
		if len(output) > 0 {
			folder = output
		}

		dscPath, err := filepath.EvalSymlinks(dscPath)
		if err != nil {
			return fmt.Errorf("failed to resolve path %s: %w", dscPath, err)
		}

		f, err := dyld.Open(dscPath)
		if err != nil {
			return err
		}
		defer f.Close()

		if addStubs {
			if len(cacheFile) == 0 {
				cacheFile = dscPath + ".a2s"
			}
			if err := f.OpenOrCreateA2SCache(cacheFile); err != nil {
				return err
			}
		}

		if dumpALL {
			// set images to all images in shared cache
			images = f.Images
			// initialize progress bar
			p = mpb.New(mpb.WithWidth(80))
			// adding a single bar, which will inherit container's width
			name := "      "
			bar = p.New(int64(len(images)),
				// progress bar filler with customized style
				mpb.BarStyle().Lbound("[").Filler("=").Tip(">").Padding("-").Rbound("|"),
				mpb.PrependDecorators(
					decor.Name(name, decor.WC{W: len(name), C: decor.DindentRight | decor.DextraSpace}),
					// replace ETA decorator with "done" message, OnComplete event
					decor.OnComplete(
						decor.AverageETA(decor.ET_STYLE_GO, decor.WC{W: 4}), "✅ ",
					),
				),
				mpb.AppendDecorators(
					decor.CountersNoUnit("%d/%d"),
					decor.Name(" ] "),
				),
			)
			log.Infof("Extracting all dylibs from %s", dscPath)
		} else {
			// get images from args
			images = make([]*dyld.CacheImage, 0, len(args)-1)
			for _, arg := range args[1:] {
				image, err := f.Image(arg)
				if err != nil {
					return err
				}
				images = append(images, image)
			}
		}

		for _, image := range images {
			m, err := image.GetMacho()
			if err != nil {
				return err
			}

			fname := filepath.Join(folder, filepath.Base(image.Name)) // default to NOT full dylib path
			if dumpALL {
				fname = filepath.Join(folder, image.Name)
			}

			if _, err := os.Stat(fname); os.IsNotExist(err) || forceExtract {
				var dcf *fixupchains.DyldChainedFixups
				if m.HasFixups() {
					dcf, err = m.DyldChainedFixups()
					if err != nil {
						log.Errorf("failed to parse fixups from in memory MachO for %s: %v", filepath.Base(image.Name), err)
					}
				}

				image.ParseLocalSymbols(false)

				syms := image.GetLocalSymbolsAsMachoSymbols()

				if addObjc && m.HasObjC() {
					log.Info("Adding ObjC symbols")
					if protos, err := m.GetObjCProtocols(); err == nil {
						for _, proto := range protos {
							syms = append(syms, macho.Symbol{
								Name:  proto.Name,
								Value: proto.Ptr,
								Desc:  0xa00,
							})
						}
					} else if !errors.Is(err, macho.ErrObjcSectionNotFound) {
						log.Error(err.Error())
					}
					if classes, err := m.GetObjCClasses(); err == nil {
						for _, class := range classes {
							syms = append(syms, macho.Symbol{
								Name:  class.Name,
								Value: class.ClassPtr,
								Desc:  0xa00,
							})
							for _, cmeth := range class.ClassMethods {
								syms = append(syms, macho.Symbol{
									Name:  fmt.Sprintf("+[%s %s]", class.Name, cmeth.Name),
									Value: cmeth.ImpVMAddr,
									Desc:  0xa00,
								})
							}
							for _, imeth := range class.InstanceMethods {
								syms = append(syms, macho.Symbol{
									Name:  fmt.Sprintf("-[%s %s]", class.Name, imeth.Name),
									Value: imeth.ImpVMAddr,
									Desc:  0xa00,
								})
							}
						}
					} else if !errors.Is(err, macho.ErrObjcSectionNotFound) {
						log.Error(err.Error())
					}
					if cats, err := m.GetObjCCategories(); err == nil {
						for _, cat := range cats {
							syms = append(syms, macho.Symbol{
								Name:  cat.Name,
								Value: cat.VMAddr,
								Desc:  0xa00,
							})
							for _, imeth := range cat.InstanceMethods {
								syms = append(syms, macho.Symbol{
									Name:  fmt.Sprintf("-[%s %s]", cat.Name, imeth.Name),
									Value: imeth.ImpVMAddr,
									Desc:  0xa00,
								})
							}
						}
					} else if !errors.Is(err, macho.ErrObjcSectionNotFound) {
						log.Error(err.Error())
					}
				}

				if addStubs {
					log.Info("Adding Stub Islands symbols")
					stubIslands, err := f.GetStubIslands()
					if err != nil {
						return err
					}
					for addr, sym := range stubIslands {
						syms = append(syms, macho.Symbol{
							Name:  sym,
							Value: addr,
							Desc:  0xa00,
						})
					}
				}

				if err := m.Export(fname, dcf, m.GetBaseAddress(), syms); err != nil {
					var perr *fs.PathError
					if errors.As(err, &perr) {
						return fmt.Errorf("failed to extract dylib %s: %v (try again with the '--output' flag to write dylib to a writable folder)", image.Name, err)
					}
					return fmt.Errorf("failed to extract dylib %s: %v", image.Name, err)
				}
				if slide {
					log.Info("Applying DSC slide-info")
					if err := rebaseMachO(f, fname); err != nil {
						return fmt.Errorf("failed to rebase dylib via cache slide info: %v", err)
					}
				}

				if dumpALL {
					bar.Increment()
				} else {
					log.Infof("Created %s", fname)
				}
			} else {
				if dumpALL {
					bar.Increment()
				} else {
					log.Warnf("Dylib already exists: %s", fname)
				}
			}

			m.Close()
		}

		if dumpALL {
			p.Wait()
		}

		return nil
	},
}
