package handlers

import (
	"archive/zip"
	"bytes"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	ahocorasick "github.com/BobuSumisu/aho-corasick"
	"github.com/avast/apkparser"
	dextk "github.com/csnewman/dextk"

	logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
	"github.com/trufflesecurity/trufflehog/v3/pkg/engine/defaults"
	"github.com/trufflesecurity/trufflehog/v3/pkg/iobuf"
)

// General Note: There are tools that can fully decompile an apk (e.g. jadx, apktool, etc.)
// However, none of these are in golang + they take awhile to run +
// they will decompile files that most likely don't contain secrets. So instead, we have a
// lightweight version that will search for secrets in the most common files that contain them.
// And run in a fraction of the time (ex: 15 seconds vs. 5 minutes)

// ToDo: Scan nested APKs (aka XAPK files). ATM the archive.go file will skip over them.
// ToDo: Provide file location information to secret output.

var (
	keywordMatcherOnce sync.Once
	keywordMatcher     *detectorKeywordMatcher
)

func defaultDetectorKeywords() []string {
	allDetectors := defaults.DefaultDetectors()

	// Remove keywords that cause lots of false positives.
	var exclusions = []string{
		"AKIA", "SG.", "pat", "token", "gh", "github", "sql", "database", "http", "key", "api-", "sdk-", "float", "-us", "gh", "pat", "token", "sid", "http", "private", "key", "segment", "close", "protocols", "verifier", "box", "privacy", "dm", "sl.", "vf", "flat",
	}

	var keywords []string
	exclusionSet := make(map[string]struct{})
	for _, excl := range exclusions {
		exclusionSet[strings.ToLower(excl)] = struct{}{}
	}

	// Aggregate all keywords from detectors.
	for _, detector := range allDetectors {
		for _, kw := range detector.Keywords() {
			kwLower := strings.ToLower(kw)
			if _, excluded := exclusionSet[kwLower]; !excluded {
				keywords = append(keywords, kwLower)
			}
		}
	}
	return keywords
}

// detectorKeywordMatcher encapsulates the Aho-Corasick trie for efficient keyword matching.
// It is used to scan APK file contents for keywords associated with our credential detectors.
// By only processing files/sections that contain these keywords, we can efficiently filter
// out irrelevant data and focus on content that is more likely to contain credentials.
// The Aho-Corasick algorithm provides fast, simultaneous matching of multiple patterns in
// a single pass through the text, which is crucial for performance when scanning large APK files.
type detectorKeywordMatcher struct{ trie *ahocorasick.Trie }

// getDefaultDetectorKeywordMatcher creates or returns the singleton detectorKeywordMatcher.
// This is implemented as a singleton for several important reasons:
// 1. Building the Aho-Corasick trie is computationally expensive and should only be done once.
// 2. The trie is immutable after construction and can be safely shared across goroutines.
// 3. The keyword list from the detectors is static for a given program execution.
// 4. Memory efficiency - we avoid duplicating the trie structure for each handler instance.
func getDefaultDetectorKeywordMatcher() *detectorKeywordMatcher {
	keywordMatcherOnce.Do(func() {
		keywords := defaultDetectorKeywords()
		keywordMatcher = &detectorKeywordMatcher{
			trie: ahocorasick.NewTrieBuilder().AddStrings(keywords).Build(),
		}
	})
	return keywordMatcher
}

// FindKeywords scans the input text and returns a slice of matched keywords.
// The method is thread-safe and uses a read lock since the trie is immutable.
// It returns unique matches only, eliminating duplicates that may occur when
// the same keyword appears multiple times in the input text.
func (km *detectorKeywordMatcher) FindKeywords(text []byte) []string {
	matches := km.trie.Match(bytes.ToLower(text))
	found := make([]string, 0, len(matches))
	seen := make(map[string]struct{}) // To avoid duplicate entries

	for _, match := range matches {
		keyword := match.MatchString()
		if _, exists := seen[keyword]; !exists {
			found = append(found, keyword)
			seen[keyword] = struct{}{}
		}
	}
	return found
}

var (
	stringInstructionType  = "const-string"
	targetInstructionTypes = []string{stringInstructionType, "iput-object", "sput-object", "const-class", "invoke-virtual", "invoke-super", "invoke-direct", "invoke-static", "invoke-interface"}
	// Note: We're only looking at a subset of instructions.
	// If expanding, update precompiled REGEX below.
	// - const-string: loads a string into a register (value)
	// - iput-object: stores a string into a field (key)
	// - the rest have to do with function, methods, objects and classes.
	reIPutRegex       = regexp.MustCompile(`iput-object obj=\d+ field=com/[a-zA-Z0-9/_]+:([a-zA-Z0-9_]+):`)
	reSPutRegex       = regexp.MustCompile(`sput-object field=com/[a-zA-Z0-9/_]+:([a-zA-Z0-9_]+):`)
	reConstRegex      = regexp.MustCompile(`const-string(?:/jumbo)? dst=\d+ value='([^']*)'`)
	reConstClassRegex = regexp.MustCompile(`const-class dst=\d+ value='[a-zA-Z0-9/_$]+/([a-zA-Z0-9]+)(?:\$|;)`)
	reInvokeRegex     = regexp.MustCompile(`invoke-(?:virtual|super|direct|static|interface)(?:/range)? method=[a-zA-Z0-9/._$]+/([a-zA-Z0-9_$]+:[a-zA-Z0-9_<]+)`)
	reInstructions    = []*regexp.Regexp{
		reIPutRegex,
		reSPutRegex,
		reConstRegex,
		reConstClassRegex,
		reInvokeRegex,
	}
)

// apkHandler handles apk archive formats.
type apkHandler struct {
	keywordMatcher *detectorKeywordMatcher
	*defaultHandler
}

// newAPKHandler creates an apkHandler.
func newAPKHandler() *apkHandler {
	return &apkHandler{
		defaultHandler: newDefaultHandler(apkHandlerType),
		keywordMatcher: getDefaultDetectorKeywordMatcher(),
	}
}

// HandleFile processes apk formatted files.
// Fatal errors that will stop processing:
// - Unable to create ZIP reader from input
// - Unable to parse resources.arsc file
// - Panics during processing (recovered but returned as errors)
//
// Non-fatal errors that will be logged and continue processing:
// - Failed to process individual files within the APK
// - Failed to process resources.arsc contents
// - Failed to process individual dex classes
// - Failed to decode specific XML files
func (h *apkHandler) HandleFile(ctx logContext.Context, input fileReader) chan DataOrErr {
	apkChan := make(chan DataOrErr, defaultBufferSize)

	go func() {
		defer close(apkChan)

		// Defer a panic recovery to handle any panics that occur during the APK processing.
		defer func() {
			if r := recover(); r != nil {
				// Return the panic as an error.
				var panicErr error
				if e, ok := r.(error); ok {
					panicErr = e
				} else {
					panicErr = fmt.Errorf("panic occurred: %v", r)
				}
				ctx.Logger().Error(panicErr, "Panic occurred when reading apk archive")
			}
		}()

		start := time.Now()

		zipReader, err := createZipReader(input)
		if err != nil {
			h.measureLatencyAndHandleErrors(ctx, start, err, apkChan)
			return
		}

		resTable, err := parseResTable(zipReader)
		if err != nil {
			h.measureLatencyAndHandleErrors(ctx, start, err, apkChan)
			return
		}

		if pkgName := extractPackageName(zipReader, resTable); pkgName != "" {
			ctx = logContext.WithValues(ctx, "apk_package", pkgName)
		}

		err = h.processAPK(ctx, zipReader, resTable, apkChan)
		if err == nil {
			h.metrics.incFilesProcessed()
		}

		h.measureLatencyAndHandleErrors(ctx, start, err, apkChan)
	}()

	return apkChan
}

// processAPK processes the apk file and sends the extracted data to the provided channel.
func (h *apkHandler) processAPK(ctx logContext.Context, zipReader *zip.Reader, resTable *apkparser.ResourceTable, apkChan chan DataOrErr) error {
	// Process the ResourceTable file for secrets
	if err := h.processResources(ctx, resTable, apkChan); err != nil {
		ctx.Logger().Error(err, "failed to process resources.arsc")
	}

	// Process all files for secrets
	for _, file := range zipReader.File {
		if err := h.processFile(ctx, file, resTable, apkChan); err != nil {
			ctx.Logger().V(2).Info("failed to process file", "file", file.Name, "error", err)
		}
	}
	return nil
}

// processResources processes the resources.arsc file and sends the extracted data to the provided channel.
func (h *apkHandler) processResources(ctx logContext.Context, resTable *apkparser.ResourceTable, apkChan chan DataOrErr) error {
	if resTable == nil {
		return errors.New("ResourceTable is nil")
	}
	provider := &apkResourceTable{table: resTable}
	rscStrRdr := extractStringsFromResTable(provider)
	return h.handleAPKFileContent(ctx, rscStrRdr, "resources.arsc", apkChan)
}

// processFile processes the file and sends the extracted data to the provided channel.
func (h *apkHandler) processFile(
	ctx logContext.Context,
	file *zip.File,
	resTable *apkparser.ResourceTable,
	apkChan chan DataOrErr,
) error {
	// check if the file is empty
	if file.UncompressedSize64 == 0 {
		return nil
	}

	// Open the file from the zip archive
	f, err := openFile(file)
	if err != nil {
		return fmt.Errorf("failed to read file %s: %w", file.Name, err)
	}
	defer func() { _ = f.Close() }()

	rdr := iobuf.NewBufferedReaderSeeker(f)
	defer func() { _ = rdr.Close() }()

	var contentReader io.Reader
	// Decode the file based on its extension
	switch strings.ToLower(filepath.Ext(file.Name)) {
	case ".xml":
		contentReader, err = decodeXML(rdr, resTable)
		if err != nil {
			return fmt.Errorf("failed to decode xml file %s: %w", file.Name, err)
		}
	case ".dex":
		contentReader, err = h.processDexFile(ctx, rdr, file.Name)
		if err != nil {
			return fmt.Errorf("failed to decode dex file %s: %w", file.Name, err)
		}
	default:
		contentReader = rdr
	}
	return h.handleAPKFileContent(ctx, contentReader, file.Name, apkChan)
}

// handleAPKFileContent sends the extracted data to the provided channel via the handleNonArchiveContent function.
func (h *apkHandler) handleAPKFileContent(
	ctx logContext.Context,
	rdr io.Reader,
	fileName string,
	apkChan chan DataOrErr,
) error {
	mimeReader, err := newMimeTypeReader(rdr)
	if err != nil {
		return fmt.Errorf("failed to create mimeTypeReader for file %s: %w", fileName, err)
	}
	ctx = logContext.WithValues(
		ctx,
		"filename", fileName,
	)
	return h.handleNonArchiveContent(ctx, mimeReader, apkChan)
}

// createZipReader creates a new ZIP reader from the input fileReader.
func createZipReader(input fileReader) (*zip.Reader, error) {
	size, err := input.Size()
	if err != nil {
		return nil, err
	}
	zipReader, err := zip.NewReader(input, size)
	if err != nil {
		return nil, err
	}
	return zipReader, err
}

// parseResTable parses the resources.arsc file and returns the ResourceTable.
func parseResTable(zipReader *zip.Reader) (*apkparser.ResourceTable, error) {
	for _, file := range zipReader.File {
		if file.Name == "resources.arsc" {
			rdr, err := openFile(file)
			if err != nil {
				return nil, err
			}

			resTable, err := apkparser.ParseResourceTable(rdr)
			_ = rdr.Close()
			if err != nil {
				return nil, err
			}
			return resTable, nil
		}
	}
	return nil, errors.New("resources.arsc file not found in the APK archive")
}

// openFile opens the file from the zip archive and returns the data as an io.ReadCloser
// Note: responsibility of calling function to close the reader
func openFile(file *zip.File) (io.ReadCloser, error) {
	rc, err := file.Open()
	if err != nil {
		return nil, err
	}
	return rc, nil
}

// resourceEntryData holds the extracted fields from a single resource table entry.
type resourceEntryData struct {
	Key          string
	ResourceType string
	Value        string
}

// GetEntryer abstracts resource table lookups for testability.
type GetEntryer interface {
	GetEntry(id uint32) (*resourceEntryData, error)
}

// apkResourceTable adapts *apkparser.ResourceTable to the GetEntryer interface.
type apkResourceTable struct {
	table *apkparser.ResourceTable
}

// GetEntry retrieves a resource entry by ID, returning nil if the entry does not exist.
func (t *apkResourceTable) GetEntry(id uint32) (*resourceEntryData, error) {
	entry, _ := t.table.GetResourceEntry(id)
	if entry == nil {
		return nil, nil
	}
	val, err := entry.GetValue().String()
	return &resourceEntryData{
		Key:          entry.Key,
		ResourceType: entry.ResourceType,
		Value:        val,
	}, err
}

// extractStringsFromResTable extracts the strings from the resources table, which is provided by the
// `GetEntryer` interface (a thin wrapper around `apkparser.ResourceTable` for testability).
// APK strings are typically stored in the 0x7f000000-0x7fffffff range.
// https://chromium.googlesource.com/chromium/src/+/master/build/android/docs/life_of_a_resource.md
func extractStringsFromResTable(provider GetEntryer) io.Reader {
	var resourceStrings bytes.Buffer
	inStrings := false
	for i := 0x7f000000; i <= 0x7fffffff; i++ {
		entry, err := provider.GetEntry(uint32(i))
		if entry == nil {
			continue
		}
		if entry.ResourceType == "string" {
			if err != nil {
				continue
			}
			inStrings = true
			resourceStrings.WriteString(entry.Key)
			resourceStrings.WriteString(": ")
			resourceStrings.WriteString(entry.Value)
			resourceStrings.WriteString("\n")
		}
		// The contiguity assumption is valid for standard AAPT output, but we might want to revisit this if there's
		// shared libs or multiple packages in a group.
		if inStrings && entry.ResourceType != "string" {
			break
		}
	}
	return &resourceStrings
}

// processDexFile decodes the dex file and returns the relevant instructions
func (h *apkHandler) processDexFile(ctx logContext.Context, rdr io.ReaderAt, dexFile string) (io.Reader, error) {
	dexReader, err := dextk.Read(rdr, dextk.WithReadCache(16))
	if err != nil {
		return nil, err
	}

	var dexOutput bytes.Buffer
	var classErrors int
	// Use an index-based loop instead of ClassIter to guarantee progress.
	// The iterator's pos field only advances after a successful `ReadClassAndParse`, so a panic would leave it stuck
	// on the same class and the HasNext loop would spin forever.
	for id := uint32(0); id < dexReader.ClassDefCount; id++ {
		func() {
			defer func() {
				if r := recover(); r != nil {
					classErrors++
				}
			}()

			node, err := dexReader.ReadClassAndParse(id)
			if err != nil {
				classErrors++
				return
			}
			h.processDexClass(ctx, dexReader, node, &dexOutput)
		}()
	}
	if classErrors > 0 {
		ctx.Logger().V(2).Info("skipped malformed dex classes", "dex_file", dexFile, "skipped_classes", classErrors)
	}

	return &dexOutput, nil
}

// processDexClass processes a single class node's methods
func (h *apkHandler) processDexClass(
	ctx logContext.Context,
	dexReader *dextk.Reader,
	node dextk.ClassNode,
	dexOutput *bytes.Buffer,
) {
	defer func() {
		if r := recover(); r != nil {
			ctx.Logger().V(2).Info("panic processing dex class", "class", node.Name.Parsed, "error", fmt.Sprintf("%v", r))
		}
	}()

	var classOutput bytes.Buffer
	methodValues := make(map[string]struct{})

	// Process Direct Methods
	processDexMethod(ctx, dexReader, node.DirectMethods, &classOutput, methodValues)
	// Process Virtual Methods
	processDexMethod(ctx, dexReader, node.VirtualMethods, &classOutput, methodValues)

	// Write the classOutput to the dexOutput
	dexOutput.Write(classOutput.Bytes())

	// Check if classOutput contains any of the default keywords
	foundKeywords := h.keywordMatcher.FindKeywords(classOutput.Bytes())

	// For each found keyword, create a keyword:value pair and append to dexOutput
	for str := range methodValues {
		for _, keyword := range foundKeywords {
			dexOutput.WriteString(keyword + ":" + str + "\n")
		}
	}
}

// processDexMethod iterates over a slice of methods, processes each method,
// handles errors, and writes the output to dexOutput.
func processDexMethod(
	ctx logContext.Context,
	dexReader *dextk.Reader,
	methods []dextk.MethodNode,
	classOutput *bytes.Buffer,
	methodValues map[string]struct{},
) {
	for _, method := range methods {
		s, err := parseDexInstructions(dexReader, method, methodValues)
		if err != nil {
			ctx.Logger().V(2).Info("failed to process dex method", "error", err)
			continue
		}
		classOutput.Write(s.Bytes())
	}
}

// parseDexInstructions processes a dex method and returns the string representation of the instruction.
func parseDexInstructions(r *dextk.Reader, m dextk.MethodNode, methodValues map[string]struct{}) (buf *bytes.Buffer, err error) {
	var instrBuf bytes.Buffer

	if m.CodeOff == 0 {
		return &instrBuf, nil
	}

	// This deferred recovery catches panics from dextk which can occur with obfuscated APKs.
	// This ensures it keeps processing and returns a valid *bytes.Buffer pointer even when a panic occurs.
	// Specifically, a panic like `runtime error: slice bounds out of range [65536:16384]` becomes a normal error
	// return. The caller (processDexMethod) already handles errors gracefully. It logs and moves on to the next
	// method. Without this recovery, the panic would propagate up and abort the entire APK.
	// Finally, it returns a valid non-zero *bytes.Buffer even when a panic occurs.
	defer func() {
		if r := recover(); r != nil {
			if e, ok := r.(error); ok {
				err = e
			} else {
				err = fmt.Errorf("%v", r)
			}
			buf = &instrBuf
		}
	}()

	c, err := r.ReadCodeAndParse(m.CodeOff)
	if err != nil {
		return &instrBuf, err
	}

	// Iterate over the instructions and extract possible secrets into `methodValues` if the instruction type is
	// `const-string`.
	for _, o := range c.Ops {
		oStr := o.String()

		instructionType := getInstructionType(oStr)
		if instructionType == "" {
			continue
		}

		val := formatAndFilterInstruction(oStr)
		if val != "" {
			instrBuf.WriteString(val + "\n")
			if instructionType == stringInstructionType {
				methodValues[val] = struct{}{}
			}
		}
	}
	return &instrBuf, nil
}

// getInstructionType checks for specific target instructions
func getInstructionType(instruction string) string {
	for _, t := range targetInstructionTypes {
		if strings.HasPrefix(instruction, t) {
			return t
		}
	}
	return ""
}

// formatAndFilterInstruction looks for a match to our regex and returns it
// Note: This is critical for ensuring secret + keyword are in close proximity.
// If we expand the instructions we're looking at, this function will need to be updated.
func formatAndFilterInstruction(line string) string {
	for _, re := range reInstructions {
		matches := re.FindStringSubmatch(line)
		if len(matches) > 1 {
			return matches[1]
		}
	}
	return ""
}

func decodeXML(rdr io.ReadSeeker, resTable *apkparser.ResourceTable) (io.Reader, error) {
	// Create a buffer to store the formatted XML data
	// Note: in the future, consider a custom writer that spills to disk if the buffer gets too large
	var buf bytes.Buffer
	enc := xml.NewEncoder(&buf)

	// Parse the XML data using the apkparser library + resource table
	err := apkparser.ParseXml(rdr, enc, resTable)
	if err == nil {
		return &buf, nil
	}

	// If the error is due to plaintext XML, return the plaintext XML.
	if errors.Is(err, apkparser.ErrPlainTextManifest) {
		if _, err := rdr.Seek(0, io.SeekStart); err != nil {
			return rdr, fmt.Errorf("error resetting reader after XML parsing error: %w", err)
		}
		return rdr, nil
	}
	return nil, err
}

// extractPackageName attempts to extract the Android package name from AndroidManifest.xml.
// Returns an empty string if the manifest is missing or the package attribute cannot be parsed.
// This is only for logging the package name, however it can be extremely difficult to debug trufflehog when running
// against gigabytes of APKs. You will want this metadata in the error messages.
func extractPackageName(zipReader *zip.Reader, resTable *apkparser.ResourceTable) (pkgName string) {
	defer func() {
		if r := recover(); r != nil {
			pkgName = ""
		}
	}()

	for _, file := range zipReader.File {
		if file.Name != "AndroidManifest.xml" {
			continue
		}
		rc, err := openFile(file)
		if err != nil {
			return ""
		}
		defer func() { _ = rc.Close() }()

		rdr := iobuf.NewBufferedReaderSeeker(rc)
		defer func() { _ = rdr.Close() }()

		decoded, err := decodeXML(rdr, resTable)
		if err != nil {
			return ""
		}

		decoder := xml.NewDecoder(decoded)
		for {
			tok, err := decoder.Token()
			if err != nil {
				break
			}
			if se, ok := tok.(xml.StartElement); ok && se.Name.Local == "manifest" {
				for _, attr := range se.Attr {
					if attr.Name.Local == "package" {
						return attr.Value
					}
				}
				break
			}
		}
		break
	}
	return ""
}