/* radare - LGPL - Copyright 2011-2026 - pancake */

#define R_LOG_ORIGIN "ragg2"

#include <r_egg.h>
#include <r_bin.h>
#include <r_cons.h>
#include <r_main.h>
#include <r_util/r_print.h>
#include <r_util.h>

typedef struct {
	RLib *l;
	REgg *e;
	// TODO flags
	// bool oneliner;
	// bool coutput;
	// bool json;
	// bool quiet;
} REggState;

typedef struct {
	const char *name;
	const char *desc;
} REggEnv;

static REggEnv env[] = {
	{ "R2_NOPLUGINS", "do not load any plugin" }
};

static void ragg_show_env(RCons *cons, bool show_desc);

static bool __lib_egg_cb(RLibPlugin *pl, void *user, void *data) {
	REggPlugin *hand = (REggPlugin *)data;
	REggState *es = (REggState *)user;
	r_libstore_add (es->e->libstore, hand);
	return true;
}

static void __load_internal_cb(void *user) {
	REggState *es = (REggState *)user;
	r_libstore_load (es->e->libstore);
}

static void __load_plugins(REggState *es) {
	es->l->cb_internal = __load_internal_cb;
	es->l->cb_internal_user = es;
	r_lib_add_handler (es->l, R_LIB_TYPE_EGG, "egg plugins", &__lib_egg_cb, NULL, es);
	r_lib_load_default_paths (es->l, R_LIB_LOAD_DEFAULT);
}

static REggState *__es_new(void) {
	REggState *es = R_NEW0 (REggState);
	es->l = r_lib_new (NULL, NULL);
	es->e = r_egg_new ();

	const bool load_plugins = !r_sys_getenv_asbool ("R2_NOPLUGINS");
	if (load_plugins) {
		__load_plugins (es);
	} else {
		r_libstore_load (es->e->libstore);
	}
	return es;
}

static void __es_free(REggState *es) {
	if (es) {
		r_egg_free (es->e);
		r_lib_free (es->l);
		free (es);
	}
}

static int usage(RCons *cons, int v) {
	r_cons_printf (cons, "Usage: ragg2 [-FOLsrxhvz] [-a arch] [-b bits] [-k os] [-o file] [-I path]\n"
	"             [-i sc] [-E enc] [-B hex] [-c k=v] [-C file] [-p pad] [-q off]\n"
	"             [-S string] [-f fmt] [-nN dword] [-dDw off:hex] [-e expr] file|f.asm|-\n");
	if (v) {
		r_cons_printf (cons,
			" -a [arch]       select architecture (x86, mips, arm, ppc)\n"
			" -b [bits]       register size (32, 64, ..)\n"
			" -B [hexpairs]   append some hexpair bytes\n"
			" -c [k=v]        set configuration options\n"
			" -C [file]       append contents of file\n"
			" -d [off:dword]  patch dword (4 bytes) at given offset\n"
			" -D [off:qword]  patch qword (8 bytes) at given offset\n"
			" -e [egg-expr]   take egg program from string instead of file\n"
			" -E [encoder]    use specific encoder. see -L\n"
			" -f [format]     output format (raw, c, pe, elf, mach0, python, javascript)\n"
			" -F              output native format (osx=mach0, linux=elf, ..)\n"
			" -h              show this help\n"
			" -H ([var])      display variable\n"
			" -i [shellcode]  include shellcode plugin, uses options. see -L\n"
			" -I [path]       add include path\n"
			" -k [os]         operating system's kernel (linux,bsd,osx,w32)\n"
			" -L              list all plugins (shellcodes and encoders)\n"
			" -n [dword]      append 32bit number (4 bytes)\n"
			" -N [dword]      append 64bit number (8 bytes)\n"
			" -o [file]       output file\n"
			" -O              use default output file (filename without extension or a.out)\n"
			" -p [padding]    add padding after compilation (padding=n10s32)\n"
			"                 ntas : begin nop, trap, 'a', sequence\n"
			"                 NTAS : same as above, but at the end\n"
			" -P [size]       prepend debruijn pattern\n"
			" -q [fragment]   debruijn pattern offset\n"
			" -r              show raw bytes instead of hexpairs\n"
			" -s              show assembler\n"
			" -S [string]     append a string\n"
			" -v              show version\n"
			" -w [off:hex]    patch hexpairs at given offset\n"
			" -x              execute\n"
			" -X [hexpairs]   execute rop chain, using the stack provided\n"
			" -z              output in C string syntax\n");
		ragg_show_env (cons, true);
	}
	return 1;
}

static void list(RCons *cons, REgg *egg) {
	RListIter *iter;
	REggPlugin *p;
	RStrBuf *sb = r_strbuf_new ("shellcodes:\n");
	r_list_foreach (egg->libstore->plugins, iter, p) {
		if (p->type == R_EGG_PLUGIN_SHELLCODE) {
			r_strbuf_appendf (sb, "%10s : %s\n", p->meta.name, p->meta.desc);
		}
	}
	r_strbuf_append (sb, "encoders:\n");
	r_list_foreach (egg->libstore->plugins, iter, p) {
		if (p->type == R_EGG_PLUGIN_ENCODER) {
			r_strbuf_appendf (sb, "%10s : %s\n", p->meta.name, p->meta.desc);
		}
	}
	r_cons_print (cons, r_strbuf_get (sb));
	r_strbuf_free (sb);
}

static bool create(RCons *cons, const char *format, const char *arch, int bits, const ut8 *code, int codelen) {
	RBin *bin = r_bin_new ();
	bool ret = false;
	r_libstore_load (bin->libstore);
	RBinArchOptions opts;
	RBuffer *b;
	r_bin_arch_options_init (&opts, arch, bits);
	b = r_bin_create (bin, format, code, codelen, NULL, 0, &opts);
	if (b) {
		ut64 blen;
		const ut8 *tmp = r_buf_data (b, &blen);
		ret = !blen || r_cons_write (cons, tmp, blen);
		if (!ret) {
			R_LOG_ERROR ("Failed to write buffer");
		}
		r_unref (b);
	} else {
		R_LOG_ERROR ("Cannot create binary for this format '%s'", format);
	}
	r_bin_free (bin);
	return ret;
}

// ragg2 also runs inside r2, so the output goes through the sandbox
static int openfile(const char *f, int x) {
	int fd = r_sandbox_open (f, O_RDWR | O_CREAT, 0644);
	if (fd == -1) {
		fd = r_sandbox_open (f, O_RDWR, 0);
		if (fd == -1) {
			return -1;
		}
	}
#if R2__UNIX__ && !__wasi__
	if (x) {
		fchmod (fd, 0755);
	}
#endif
#if _MSC_VER || R2__WINDOWS__
	int r = _chsize (fd, 0);
#else
	int r = ftruncate (fd, 0);
#endif
	if (r != 0) {
		R_LOG_ERROR ("Could not resize");
	}
	return fd;
}
#define ISEXEC (fmt != 'r')

static void ragg_env_print(RCons *cons, const char *name) {
	char *value = r_sys_getenv (name);
	r_cons_printf (cons, "%s\n", R_STR_ISNOTEMPTY (value)? value: "");
	free (value);
}

static void ragg_show_env(RCons *cons, bool show_desc) {
	int id = 0;
	for (id = 0; id < (sizeof (env) / sizeof (env[0])); id++) {
		if (show_desc) {
			r_cons_printf (cons, "%s\t%s\n", env[id].name, env[id].desc);
		} else {
			r_cons_printf (cons, "%s=", env[id].name);
			ragg_env_print (cons, env[id].name);
		}
	}
}

static void ragg_patch_word(REgg *egg, const char *arg, int size) {
	const char *colon = strchr (arg, ':');
	if (!colon) {
		R_LOG_ERROR ("Missing colon in -%c", size == 4? 'd': 'D');
		return;
	}
	char *offstr = r_str_ndup (arg, colon - arg);
	ut64 off = r_num_math (NULL, offstr);
	ut64 n = r_num_math (NULL, colon + 1);
	free (offstr);
	// TODO: support big endian
	ut8 word[8];
	if (size == 4) {
		r_write_le32 (word, (ut32)n);
	} else {
		r_write_le64 (word, n);
	}
	r_egg_patch (egg, off, word, size);
}

R_API int r_main_ragg2(int argc, const char **argv) {
	bool own_cons;
	RCons *cons = r_main_cons_init (&own_cons);
	const char *file = NULL;
	const char *padding = NULL;
	const char *pattern = NULL;
	const char *str = NULL;
	char *bytes = NULL;
	const char *contents = NULL;
	const char *arch = R_SYS_ARCH;
	const char *os = R_EGG_OS_NAME;
	const char *format = "raw";
	bool show_execute = false;
	bool show_execute_rop = false;
	bool show_hex = true;
	bool show_asm = false;
	bool show_raw = false;
	int append = 0;
	int show_str = 0;
	ut64 get_offset = 0;
	const char *shellcode = NULL;
	const char *encoder = NULL;
	const char *eggprg = NULL;
	char *sequence = NULL;
	int bits = R_SYS_BITS_CHECK (R_SYS_BITS, 64)? 64: 32;
	int fmt = 0;
	const char *ofile = NULL;
	int ofileauto = 0;
	RBuffer *b;
	int c, i, fd = -1, ret = 0;
	bool file_context = false;

	if (argc < 2) {
		return r_main_cons_fini (cons, own_cons, usage (cons, 1));
	}

	char *include_path = r_sys_getenv (R_EGG_INCDIR_ENV);
	REggState *es = __es_new ();

	RGetopt opt;
	r_getopt_init (&opt, argc, argv, "a:b:B:c:C:d:D:e:E:f:FhH:i:I:k:Ln:N:o:Op:P:q:rsS:vw:xX:z");
	if (argc == 2 && !strcmp (argv[1], "-H")) {
		ragg_show_env (cons, false);
		goto beach;
	}
	while ((c = r_getopt_next (&opt)) != -1) {
		switch (c) {
		case 'a':
			arch = opt.arg;
			if (!strcmp (arch, "trace")) {
				show_asm = true;
				show_hex = false;
			}
			break;
		case 'e':
			eggprg = opt.arg;
			break;
		case 'E':
			encoder = opt.arg;
			break;
		case 'b':
			bits = atoi (opt.arg);
			break;
		case 'B':
			bytes = r_str_append (bytes, opt.arg);
			break;
		case 'C':
			if (R_STR_ISEMPTY (opt.arg)) {
				R_LOG_ERROR ("Cannot open empty contents path");
				goto fail;
			}
			contents = opt.arg;
			break;
		case 'w':
			{
				char *arg = strdup (opt.arg);
				char *p = strchr (arg, ':');
				if (p) {
					int len, off;
					ut8 *b;
					*p++ = 0;
					off = r_num_math (NULL, arg);
					b = calloc (1, strlen (opt.arg) + 1);
					len = r_hex_str2bin (p, b);
					if (len > 0) {
						r_egg_patch (es->e, off, (const ut8 *)b, len);
					} else {
						R_LOG_ERROR ("Invalid hexstr for -w");
					}
					free (b);
				} else {
					R_LOG_ERROR ("Missing colon in -w");
				}
				free (arg);
			}
			break;
		case 'n':
			{
				ut32 n = r_num_math (NULL, opt.arg);
				append = 1;
				r_egg_patch (es->e, -1, (const ut8 *)&n, 4);
			}
			break;
		case 'N':
			{
				ut64 n = r_num_math (NULL, opt.arg);
				r_egg_patch (es->e, -1, (const ut8 *)&n, 8);
				append = 1;
			}
			break;
		case 'd':
			ragg_patch_word (es->e, opt.arg, 4);
			break;
		case 'D':
			ragg_patch_word (es->e, opt.arg, 8);
			break;
		case 'S':
			str = opt.arg;
			break;
		case 'o':
			ofile = opt.arg;
			break;
		case 'O':
			ofileauto = 1;
			break;
		case 'I':
			if (R_STR_ISEMPTY (opt.arg)) {
				R_LOG_ERROR ("Cannot open empty include path");
				goto fail;
			}
			r_egg_lang_include_path (es->e, opt.arg);
			break;
		case 'i':
			shellcode = opt.arg;
			break;
		case 'p':
			padding = opt.arg;
			break;
		case 'P':
			pattern = opt.arg;
			break;
		case 'c':
			{
				const char *eq = strchr (opt.arg, '=');
				if (eq) {
					char *key = r_str_ndup (opt.arg, eq - opt.arg);
					r_egg_option_set (es->e, key, eq + 1);
					free (key);
				} else {
					r_egg_option_set (es->e, opt.arg, "true");
				}
			}
			break;
		case 'F':
#if __APPLE__
			format = "mach0";
#elif R2__WINDOWS__
			format = "pe";
#else
			format = "elf";
#endif
			show_asm = false;
			break;
		case 'f':
			format = opt.arg;
			show_asm = false;
			break;
		case 's':
			show_asm = true;
			show_hex = false;
			break;
		case 'k':
			os = opt.arg;
			break;
		case 'r':
			show_raw = true;
			break;
		case 'x':
			// execute
			show_execute = true;
			break;
		case 'X':
			// execute rop chain
			show_execute = 1;
			show_execute_rop = 1;
			break;
		case 'L':
			list (cons, es->e);
			goto beach;
		case 'h':
			ret = usage (cons, 1);
			goto beach;
		case 'H':
			ragg_env_print (cons, opt.arg);
			goto beach;
		case 'v':
			ret = r_main_version_print (cons, "ragg2", 0);
			goto beach;
		case 'z':
			show_str = 1;
			break;
		case 'q':
			get_offset = 1;
			free (sequence);
			sequence = strdup (opt.arg);
			break;
		default:
			goto fail;
		}
	}

	if (opt.ind == argc && !eggprg && !shellcode && !bytes && !contents && !encoder && !padding && !pattern && !append && !get_offset && !str) {
		ret = usage (cons, 0);
		goto beach;
	}
	if (opt.ind != argc) {
		file = argv[opt.ind];
	}

	if (bits == 64) {
		if (!strcmp (format, "mach0")) {
			format = "mach064";
		} else if (!strcmp (format, "elf")) {
			format = "elf64";
		} else if (!strcmp (format, "pe")) {
			format = "pe64";
		}
	}

	// catch this first
	if (get_offset) {
		if (strncmp (sequence, "0x", 2)) {
			R_LOG_ERROR ("Need hex value with `0x' prefix e.g. 0x41414142");
			goto fail;
		}

		get_offset = r_num_math (0, sequence);
		r_cons_printf (cons, "Little endian: %d\n", r_debruijn_offset (get_offset, false));
		r_cons_printf (cons, "Big endian: %d\n", r_debruijn_offset (get_offset, true));
		goto beach;
	}

	// PPC32 and PPC64 ELFv1 are big-endian by convention. ppc64le will need
	// explicit endian selection once RArchConfig integration lands (egg.c TODO).
	const int egg_endian = (arch && !strcmp (arch, "ppc"))? 1: 0;
	r_egg_setup (es->e, arch, bits, egg_endian, os);
	if (file) {
		if (R_STR_ISEMPTY (file)) {
			R_LOG_ERROR ("Cannot open empty path");
			goto fail;
		}
		if (!strcmp (file, "-")) {
			char buf[1024];
			for (;;) {
				if (!fgets (buf, sizeof (buf), stdin)) {
					break;
				}
				if (feof (stdin)) {
					break;
				}
				r_egg_load (es->e, buf, 0);
			}
		} else if (strstr (file, ".c")) {
			char *fileSanitized = strdup (file);
			r_str_sanitize (fileSanitized);
			char *textFile = r_egg_cfile_parser (fileSanitized, arch, os, bits);

			if (!textFile) {
				R_LOG_ERROR ("Failure while parsing '%s'", fileSanitized);
				goto fail;
			}

			size_t l;
			char *buf = r_file_slurp (textFile, &l);
			if (buf && l > 0) {
				r_egg_raw (es->e, (const ut8 *)buf, (int)l);
			} else {
				R_LOG_ERROR ("Cannot load '%s'", textFile);
			}

			r_file_rm (textFile);
			free (fileSanitized);
			free (textFile);
			free (buf);
		} else {
			if (strstr (file, ".s") || strstr (file, ".asm")) {
				fmt = 'a';
			} else if (strstr (file, ".bin") || strstr (file, ".raw")) {
				fmt = 'r';
			} else {
				fmt = 0;
			}
			if (!r_egg_include (es->e, file, fmt)) {
				R_LOG_ERROR ("Cannot open '%s'", file);
				goto fail;
			}
		}
	} else {
		if (eggprg && !r_egg_include_str (es->e, eggprg)) {
			R_LOG_ERROR ("Cannot parse egg program");
			goto fail;
		}
	}

	// compile source code to assembly
	if (!r_egg_compile (es->e)) {
		if (!fmt) {
			R_LOG_ERROR ("r_egg_compile failed");
			goto fail;
		}
	}

	// append the provided string
	if (str) {
		int l = strlen (str);
		if (l > 0) {
			r_egg_raw (es->e, (const ut8 *)str, l);
		}
	}

	// add raw file
	if (contents) {
		size_t l;
		char *buf = r_file_slurp (contents, &l);
		if (buf && l > 0) {
			r_egg_raw (es->e, (const ut8 *)buf, (int)l);
		} else {
			R_LOG_ERROR ("Cannot load '%s'", contents);
		}
		free (buf);
	}

	// add shellcode
	if (shellcode) {
		if (!r_egg_shellcode (es->e, shellcode)) {
			R_LOG_ERROR ("Unknown shellcode '%s'", shellcode);
			goto fail;
		}
	}

	// add raw bytes
	if (bytes) {
		ut8 *b = calloc (1, strlen (bytes) + 1);
		int len = r_hex_str2bin (bytes, b);
		if (len > 0) {
			if (!r_egg_raw (es->e, b, len)) {
				R_LOG_ERROR ("Unknown '%s'", shellcode);
				free (b);
				goto fail;
			}
		} else {
			R_LOG_ERROR ("Invalid hexpair string for -B");
		}
		free (b);
		free (bytes);
		bytes = NULL;
	}

	// Create output files without redirecting the caller's stdout.
	if (ofileauto) {
		if (file) {
			char *o, *q, *p = strdup (file);
			if ((o = strchr (p, '.'))) {
				while ((q = strchr (o + 1, '.'))) {
					o = q;
				}
				*o = 0;
				fd = openfile (p, ISEXEC);
			} else {
				fd = openfile ("a.out", ISEXEC);
			}
			free (p);
		} else {
			fd = openfile ("a.out", ISEXEC);
		}
		if (fd == -1) {
			R_LOG_ERROR ("cannot open file '%s'", opt.arg);
			goto fail;
		}
	}
	if (ofile) {
		if (fd != -1) {
			close (fd);
		}
		fd = openfile (ofile, ISEXEC);
		if (fd == -1) {
			R_LOG_ERROR ("cannot open file '%s'", ofile);
			goto fail;
		}
	}
	if (fd != -1) {
		r_cons_push (cons);
		file_context = true;
	}

	// assemble to binary
	if (!show_asm) {
		if (!r_egg_assemble (es->e)) {
			R_LOG_ERROR ("r_egg_assemble: invalid assembly");
			goto fail;
		}
	}
	if (encoder) {
		if (!r_egg_encode (es->e, encoder)) {
			R_LOG_ERROR ("Invalid encoder '%s'", encoder);
			goto fail;
		}
	}

	// add padding
	if (padding) {
		r_egg_padding (es->e, padding);
	}

	// add pattern
	if (pattern) {
		r_egg_pattern (es->e, r_num_math (NULL, pattern));
	}

	// apply patches
	if (!es->e->bin) {
		es->e->bin = r_buf_new ();
	}
	if (! (b = r_egg_get_bin (es->e))) {
		R_LOG_ERROR ("r_egg_get_bin: invalid egg :(");
		goto fail;
	}
	r_egg_finalize (es->e);

	if (show_asm) {
		char *assembly = r_egg_get_assembly (es->e);
		r_cons_println (cons, assembly);
		free (assembly);
	}

	if (show_raw || show_hex || show_execute) {
		if (show_execute) {
			if (show_execute_rop) {
				ret = r_egg_run_rop (es->e);
			} else {
				ret = r_egg_run (es->e);
			}
			goto beach;
		}
		b = r_egg_get_bin (es->e);
		if (show_raw) {
			ut64 blen;
			const ut8 *tmp = r_buf_data (b, &blen);
			if (blen && !r_cons_write (cons, tmp, blen)) {
				R_LOG_ERROR ("Failed to write buffer");
				goto fail;
			}
		} else {
			if (!format) {
				R_LOG_ERROR ("No format specified wtf");
				goto fail;
			}
			RPrint *p = r_print_new ();
			r_cons_bind (cons, &p->consb);
			ut64 tmpsz;
			const ut8 *tmp = r_buf_data (b, &tmpsz);
			switch (*format) {
			case 'c':
				r_print_code (p, 0, tmp, tmpsz, 'c');
				break;
			case 'j': // JavaScript
				r_print_code (p, 0, tmp, tmpsz, 'j');
				break;
			case 'r':
				if (show_str) {
					r_cons_printf (cons, "\"");
					for (i = 0; i < tmpsz; i++) {
						r_cons_printf (cons, "\\x%02x", tmp[i]);
					}
					r_cons_printf (cons, "\"\n");
				} else if (show_hex) {
					r_buf_seek (b, 0, R_BUF_SET);
					for (i = 0; i < tmpsz; i++) {
						r_cons_printf (cons, "%02x", tmp[i]);
					}
					r_cons_printf (cons, "\n");
				} // else show_raw is_above ()
				break;
			case 'p': // PE/python
				if (strlen (format) > 2 && format[1] == 'y') { // Python
					r_print_code (p, 0, tmp, tmpsz, 'p');
				} else { // PE
					ret = !create (cons, format, arch, bits, tmp, tmpsz);
				}
				break;
			case 'e': // ELF
			case 'm': // MACH0
				ret = !create (cons, format, arch, bits, tmp, tmpsz);
				break;
			default:
				R_LOG_ERROR ("unknown executable format (%s)", format);
				ret = 1;
				break;
			}
			r_print_free (p);
		}
	}
	goto beach;
fail:
	ret = 1;
beach:
	if (file_context) {
		size_t size;
		char *output = r_cons_drain (cons, &size);
		if (size && write (fd, output, size) != size) {
			R_LOG_ERROR ("Failed to write output file");
			ret = 1;
		}
		free (output);
		r_cons_pop (cons);
	}
	if (fd != -1) {
		close (fd);
	}
	free (bytes);
	free (sequence);
	__es_free (es);
	r_sys_setenv (R_EGG_INCDIR_ENV, include_path);
	free (include_path);
	return r_main_cons_fini (cons, own_cons, ret);
}