// SPDX-FileCopyrightText: 2023 RizinOrg <info@rizin.re>
// SPDX-FileCopyrightText: 2023 deroad <wargio@libero.it>
// SPDX-License-Identifier: LGPL-3.0-only

#include <rz_bin.h>
#include "i/private.h"

RZ_IPI const RzDemanglerPlugin *rz_bin_process_get_demangler_plugin_from_lang(RzBin *bin, RzBinLanguage language) {
	language = RZ_BIN_LANGUAGE_MASK(language);
	const char *lang_s = rz_bin_language_to_string(language);
	if (!lang_s) {
		return NULL;
	}
	return rz_demangler_plugin_get(bin->demangler, lang_s);
}

RZ_IPI void rz_bin_process_rust(RzBinObject *o, char *demangled, ut64 paddr, ut64 vaddr, bool is_method) {
	// find where to split symbol.
	char *str = demangled;
	char *ptr = NULL;
	char *name = NULL;
	while ((ptr = strstr(str, "::"))) {
		name = ptr;
		str = ptr + 2;
	}

	if (!name || RZ_STR_ISEMPTY(name + 2)) {
		return;
	}

	*name = 0;
	if (is_method) {
		rz_bin_object_add_method(o, demangled, name + 2, paddr, vaddr);
	} else {
		rz_bin_object_add_field(o, demangled, name + 2, paddr, vaddr);
	}
	*name = ':';
}

static char *find_cxx_name(char *start, const char *end) {
	char *kname = NULL;
	unsigned int tmpl = 0;
	while (start < end) {
		if (tmpl > 0) {
			if (start[0] == '<') {
				tmpl++;
			} else if (start[0] == '>') {
				tmpl--;
			}
		} else {
			if (start[0] == '<') {
				tmpl++;
			} else if (start[0] == ':' && start[1] == ':') {
				kname = start;
			}
		}
		start++;
	}
	return kname;
}

RZ_IPI void rz_bin_process_cxx(RzBinObject *o, char *demangled, ut64 paddr, ut64 vaddr) {
	if (strstr(demangled, " for ") ||
		strstr(demangled, " to ")) {
		/* these symbols are not fields nor methods. */
		return;
	}
	bool is_method = true;
	const char *limit = NULL;
	if (rz_str_startswith(demangled, "(anonymous namespace)::")) {
		limit = strchr(demangled + strlen("(anonymous namespace)::"), '(');
	} else {
		limit = strchr(demangled, '(');
	}
	if (!limit) {
		limit = demangled + (strlen(demangled) - 1);
		is_method = false;
	}

	// find where to split symbol.
	char *name = find_cxx_name(demangled, limit);
	if (!name || RZ_STR_ISEMPTY(name + 2)) {
		// eprintf("bad '%s' '%s'\n", demangled, name);
		return;
	}

	*name = 0;
	if (is_method) {
		rz_bin_object_add_method(o, demangled, name + 2, paddr, vaddr);
	} else {
		rz_bin_object_add_field(o, demangled, name + 2, paddr, vaddr);
	}
	*name = ':';
}

static char *find_swift_methodname(RZ_NONNULL char *demangled) {
	// methods can be main.Tost.deinit or main.Tost.init() -> main.Tost
	// so we will return after second dot
	char *dot = strchr(demangled, '.');
	dot = dot ? strchr(dot + 1, '.') : NULL;
	if (!dot) {
		return NULL;
	}
	char *methodname = dot + 1;
	if (RZ_STR_ISEMPTY(methodname)) {
		return NULL;
	}
	return methodname;
}

static void bin_process_metaclass(RzBinObject *o, RzBinSymbol *symbol) {
	if (RZ_STR_ISEMPTY(symbol->dname)) {
		return;
	}
	char *no_classname = strstr(symbol->dname, "full type metadata for ");
	char *classname = strstr(symbol->dname, "type metadata for ");
	if (!classname || no_classname) { // only for "type metadata for class"
		return;
	}
	rz_bin_object_add_field(o, classname + strlen("type metadata for "), symbol->dname, symbol->paddr, symbol->vaddr);
}

// This function is used to process Swift methods.
static void bin_process_swift_class_method(RzBinObject *o, RzBinSymbol *symbol) {
	// Before the second dot, we have the class name
	char *dot = strchr(symbol->dname, '.');
	dot = dot ? strchr(dot + 1, '.') : NULL;
	if (!dot) {
		return;
	}
	char *classname = rz_str_ndup(symbol->dname, dot - symbol->dname);
	// classname should not have any spaces or ( or )
	if (strchr(classname, ' ') || strchr(classname, '(') || strchr(classname, ')')) {
		free(classname);
		return;
	}
	symbol->classname = classname;
	char *methodname = find_swift_methodname(symbol->dname);

	if (!methodname) {
		free(classname);
		return;
	}

	rz_bin_object_add_class(o, classname, NULL, UT64_MAX);
	rz_bin_object_add_method(o, classname, methodname, symbol->paddr, symbol->vaddr);
}

// this process function does not work with the Apple demangler.
RZ_IPI void rz_bin_process_swift(RzBinObject *o, RzBinSymbol *symbol) {
	if (RZ_STR_ISEMPTY(symbol->dname)) {
		return;
	}
	bin_process_metaclass(o, symbol);
	bin_process_swift_class_method(o, symbol);
}

static inline bool is_known_canary(const char *name) {
	return strstr(name, "__stack_chk_fail") ||
		strstr(name, "__stack_chk_guard") ||
		strstr(name, "__intel_security_cookie") ||
		strstr(name, "__stack_smash_handler") ||
		strstr(name, "__security_init_cookie");
}

#define DEFINE_IS_SANITIZER(san_name) \
	static inline bool is_##san_name(const char *var) { \
		return strstr(var, "__" RZ_STR_DEF(san_name) "_") || \
			strstr(var, "__" RZ_STR_DEF(san_name) "::"); \
	}

DEFINE_IS_SANITIZER(sanitizer)
DEFINE_IS_SANITIZER(ubsan)
DEFINE_IS_SANITIZER(asan)
DEFINE_IS_SANITIZER(msan)
DEFINE_IS_SANITIZER(tsan)

static inline bool is_known_objc_arc(const char *name) {
	return strstr(name, "objc_retain") ||
		strstr(name, "objc_release") ||
		strstr(name, "objc_storeStrong");
}

static inline bool is_known_fortify_source(const char *name) {
	return strstr(name, "memcpy_chk") ||
		strstr(name, "memmove_chk") ||
		strstr(name, "memset_chk") ||
		strstr(name, "stpcpy_chk") ||
		strstr(name, "stpncpy_chk") ||
		strstr(name, "strcat_chk") ||
		strstr(name, "strcpy_chk") ||
		strstr(name, "strncat_chk") ||
		strstr(name, "strncpy_chk") ||
		strstr(name, "snprintf_chk") ||
		strstr(name, "sprintf_chk") ||
		strstr(name, "vsnprintf_chk") ||
		strstr(name, "vsprintf_chk");
}

static inline void bin_info_auto_detect(RzBinInfo *info, const char *name) {
	if (RZ_STR_ISEMPTY(name)) {
		return;
	} else if (is_known_canary(name)) {
		info->has_canary = true;
	} else if (is_ubsan(name)) {
		info->sanitizers |= RZ_BIN_SANITIZER_UBSAN;
	} else if (is_asan(name)) {
		info->sanitizers |= RZ_BIN_SANITIZER_ASAN;
	} else if (is_tsan(name)) {
		info->sanitizers |= RZ_BIN_SANITIZER_TSAN;
	} else if (is_msan(name)) {
		info->sanitizers |= RZ_BIN_SANITIZER_MSAN;
	} else if (is_sanitizer(name)) {
		info->sanitizers |= RZ_BIN_SANITIZER_GENERIC;
	} else if (is_known_objc_arc(name)) {
		info->has_objc_arc = true;
	} else if (is_known_fortify_source(name)) {
		info->has_fortify_source = true;
	}
}

RZ_IPI void rz_bin_info_auto_resolve_fields(RzBinObject *o) {
	void **iter;
	RzBinImport *import;
	RzBinSymbol *symbol;
	RzBinInfo *info = o->info;
	if (!info) {
		return;
	}

	info->sanitizers = RZ_BIN_SANITIZER_NONE;
	rz_pvector_foreach (o->imports, iter) {
		import = *iter;
		bin_info_auto_detect(info, import->name);
	}

	rz_pvector_foreach (o->symbols, iter) {
		symbol = *iter;
		bin_info_auto_detect(info, symbol->name);
	}
}

/**
 * \brief      Initialize the data of the given RzBinObject using the defined RzBinPlugin
 *
 * \param      bf    The RzBinFile to use
 * \param      o     The RzBinObject to initialize
 *
 * \return     On success returns true
 */
RZ_IPI bool rz_bin_object_process_plugin_data(RZ_NONNULL RzBinFile *bf, RZ_NONNULL RzBinObject *o) {
	rz_return_val_if_fail(bf && bf->rbin && o && o->plugin, false);
	const RzDemanglerPlugin *demangler = NULL;

	rz_bin_set_and_process_file(bf, o);
	rz_bin_set_and_process_entries(bf, o);
	rz_bin_set_and_process_maps(bf, o);
	rz_bin_set_imports_from_plugin(bf, o);
	rz_bin_set_symbols_from_plugin(bf, o);
	rz_bin_set_and_process_sections(bf, o);
	rz_bin_set_and_process_fields(bf, o);
	rz_bin_set_and_process_classes(bf, o);

	// we need to detect the language of the binary
	// one way can be based on the compiler.
	if (o->info) {
		char *go_compiler = rz_bin_file_golang_compiler(bf);
		if (go_compiler) {
			o->info->lang = "go";
			o->lang = RZ_BIN_LANGUAGE_GO;
			if (RZ_STR_ISNOTEMPTY(o->info->compiler)) {
				char *merge = rz_str_newf("%s %s", go_compiler, o->info->compiler);
				free(o->info->compiler);
				free(go_compiler);
				o->info->compiler = merge;
			} else {
				free(o->info->compiler);
				o->info->compiler = go_compiler;
			}
		}
	}

	// or based on the fetched data in the RzBinFile
	if (RZ_BIN_LANGUAGE_MASK(o->lang) == RZ_BIN_LANGUAGE_UNKNOWN) {
		o->lang = rz_bin_language_detect(bf);
	}
	// Process strings after the language was set,
	// because some languages imply a specific encoding.
	rz_bin_set_and_process_strings(bf, o);

	// now we can process the data.
	RzDemanglerFlag flags = rz_demangler_get_flags(bf->rbin->demangler);
	if (bf->rbin->demangle) {
		demangler = rz_bin_process_get_demangler_plugin_from_lang(bf->rbin, o->lang);
	}
	rz_bin_process_symbols(bf, o, demangler, flags);
	rz_bin_process_imports(bf, o, demangler, flags);
	rz_bin_set_and_process_relocs(bf, o, demangler, flags);
	rz_bin_info_auto_resolve_fields(o);

	return true;
}

/**
 * \brief Remove all previously identified strings in the binary object and scan it again for strings.
 */
RZ_API bool rz_bin_object_reset_strings(RZ_NONNULL RzBin *bin, RZ_NONNULL RzBinFile *bf, RZ_NONNULL RzBinObject *obj) {
	rz_return_val_if_fail(bin && bf && obj, false);
	rz_bin_set_and_process_strings(bf, obj);
	return obj->strings != NULL;
}

/**
 * \brief      If the given flags differs from the one already set, then re-demangles all the symbols.
 *
 * \param      bin    The RzBin context to used for demangling
 * \param[in]  flags  The demangler flags to apply
 */
RZ_API void rz_bin_demangle_with_flags(RZ_NONNULL RzBin *bin, RzDemanglerFlag flags) {
	rz_return_if_fail(bin && bin->binfiles);

	RzDemanglerFlag f = rz_demangler_get_flags(bin->demangler);
	if ((f & flags) == flags) {
		return;
	}
	rz_demangler_set_flags(bin->demangler, flags);

	if (!bin->demangle) {
		return;
	}

	RzBinFile *bf = NULL;
	RzListIter *it = NULL;
	// reload each bins and demangle properly
	rz_list_foreach (bin->binfiles, it, bf) {
		if (!bf || !bf->o) {
			continue;
		}
		RzBinObject *o = bf->o;
		const RzDemanglerPlugin *demangler = rz_bin_process_get_demangler_plugin_from_lang(bin, o->lang);
		rz_bin_demangle_relocs_with_flags(o, demangler, flags);
		rz_bin_demangle_imports_with_flags(o, demangler, flags);
		rz_bin_demangle_symbols_with_flags(o, demangler, flags);
	}
}
