//@category Analysis
import java.io.*;
import java.util.*;

import ghidra.app.decompiler.*;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.listing.*;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.symbol.*;

public class DumpStealerCheck extends GhidraScript {

	private static final String[] API_NEEDLES = {
		"WinHttp", "InternetOpen", "InternetConnect", "HttpSend", "InternetRead",
		"URLDownload", "URLOpen", "CryptUnprotect", "CryptProtect",
		"CreateRemoteThread", "WriteProcessMemory", "ReadProcessMemory",
		"VirtualAlloc", "VirtualProtect", "NtWriteVirtual", "NtReadVirtual",
		"OpenProcess", "SetWindowsHook", "GetAsyncKeyState", "GetKeyState",
		"GetKeyboardState", "BitBlt", "PrintWindow", "GetClipboard",
		"FtpPut", "FtpOpen", "WSASend", "send", "connect", "InternetWrite",
		"HttpOpenRequest", "WinHttpSend", "WinHttpRead", "WinHttpConnect",
		"ShellExecute", "WinExec", "CreateProcess", "RegSetValue",
		"InternetGetCookie", "InternetSetCookie", "CredEnumerate", "CredRead"
	};

	private static final String[] STR_NEEDLES = {
		"Login Data", "Web Data", "User Data\\Default", "Google\\Chrome",
		"Microsoft\\Edge", "discord.com/api/webhooks", "api.telegram.org",
		"wallet.dat", "Exodus", "Electrum", "MetaMask", "FileZilla",
		"KeePass", "\\tdata", "stealer", "grabber", "clipper",
		"CryptUnprotectData", "leveldb", "loginusers.vdf", "1Password"
	};

	@Override
	public void run() throws Exception {
		String outDir = "C:\\Users\\Admin\\AppData\\Local\\Temp\\jx_ghidra_out";
		if (getScriptArgs() != null && getScriptArgs().length > 0) {
			outDir = getScriptArgs()[0];
		}
		new File(outDir).mkdirs();
		String name = currentProgram.getName().replaceAll("[^A-Za-z0-9._-]", "_");
		File out = new File(outDir, name + ".decomp.txt");
		PrintWriter pw = new PrintWriter(new OutputStreamWriter(new FileOutputStream(out), "UTF-8"));

		try {
			pw.println("PROGRAM=" + currentProgram.getName());
			pw.println("PATH=" + currentProgram.getExecutablePath());
			pw.println("LANG=" + currentProgram.getLanguageID());
			pw.println("IMAGE_BASE=" + currentProgram.getImageBase());
			pw.println();

			pw.println("===== MEMORY BLOCKS =====");
			for (MemoryBlock block : currentProgram.getMemory().getBlocks()) {
				pw.printf("%s va=%s size=0x%x init=%s perm=%s%n",
					block.getName(), block.getStart(), block.getSize(),
					block.isInitialized(),
					(block.isRead() ? "R" : "") + (block.isWrite() ? "W" : "") + (block.isExecute() ? "X" : ""));
			}

			pw.println();
			pw.println("===== IMPORTS =====");
			SymbolTable st = currentProgram.getSymbolTable();
			List<String> imports = new ArrayList<String>();
			for (Symbol sym : st.getExternalSymbols()) {
				imports.add(sym.getParentNamespace().getName() + "!" + sym.getName());
			}
			Collections.sort(imports);
			for (String imp : imports) {
				pw.println(imp);
			}

			pw.println();
			pw.println("===== DEFINED STRING HITS =====");
			Listing listing = currentProgram.getListing();
			DataIterator dit = listing.getDefinedData(true);
			int strHits = 0;
			while (dit.hasNext() && !monitor.isCancelled()) {
				Data d = dit.next();
				if (!d.hasStringValue()) {
					continue;
				}
				String s = d.getDefaultValueRepresentation();
				if (s == null) {
					continue;
				}
				String lower = s.toLowerCase();
				for (String n : STR_NEEDLES) {
					if (lower.contains(n.toLowerCase())) {
						pw.println(d.getAddress() + " " + s);
						strHits++;
						break;
					}
				}
			}
			if (strHits == 0) {
				pw.println("(none)");
			}

			Set<Function> targets = new LinkedHashSet<Function>();
			pw.println();
			pw.println("===== INTERESTING API XREFS =====");
			for (Symbol sym : st.getExternalSymbols()) {
				String nm = sym.getName();
				boolean interesting = false;
				for (String n : API_NEEDLES) {
					if (nm.contains(n)) {
						interesting = true;
						break;
					}
				}
				if (!interesting) {
					continue;
				}
				Reference[] refs = getReferencesTo(sym.getAddress());
				if (refs == null || refs.length == 0) {
					// thunk / external loc
					for (Symbol alias : st.getSymbols(sym.getAddress())) {
						refs = getReferencesTo(alias.getAddress());
						if (refs != null && refs.length > 0) {
							break;
						}
					}
				}
				pw.println("API " + sym.getParentNamespace().getName() + "!" + nm + " refs=" +
					(refs == null ? 0 : refs.length));
				if (refs == null) {
					continue;
				}
				for (Reference ref : refs) {
					Function f = listing.getFunctionContaining(ref.getFromAddress());
					if (f != null) {
						pw.println("  from " + f.getName() + " @ " + f.getEntryPoint() +
							" callsite=" + ref.getFromAddress());
						targets.add(f);
					}
					else {
						pw.println("  from " + ref.getFromAddress() + " (no function)");
					}
				}
			}

			Function entry = listing.getFunctionAt(currentProgram.getAddressFactory()
				.getDefaultAddressSpace()
				.getAddress(currentProgram.getImageBase().getOffset()));
			Symbol entrySym = st.getPrimarySymbol(currentProgram.getAddressFactory()
				.getDefaultAddressSpace()
				.getAddress(0));
			Function ep = currentProgram.getListing().getFunctionAt(
				currentProgram.getImageBase());
			if (ep != null) {
				targets.add(ep);
			}
			FunctionIterator fi = listing.getFunctions(true);
			while (fi.hasNext()) {
				Function f = fi.next();
				String n = f.getName().toLowerCase();
				if (n.contains("winmain") || n.equals("entry") || n.contains("maincrt") ||
					n.equals("_main") || n.equals("main")) {
					targets.add(f);
				}
			}

			pw.println();
			pw.println("===== DECOMPILED FUNCTIONS (" + targets.size() + ") =====");
			DecompInterface decomp = new DecompInterface();
			DecompileOptions options = new DecompileOptions();
			decomp.setOptions(options);
			decomp.toggleCCode(true);
			decomp.toggleSyntaxTree(false);
			decomp.setSimplificationStyle("decompile");
			if (!decomp.openProgram(currentProgram)) {
				pw.println("Decompiler open failed: " + decomp.getLastMessage());
				return;
			}
			try {
				int count = 0;
				for (Function f : targets) {
					if (monitor.isCancelled()) {
						break;
					}
					if (count >= 40) {
						pw.println("(truncated after 40 functions)");
						break;
					}
					DecompileResults res = decomp.decompileFunction(f, 60, monitor);
					pw.println();
					pw.println("----- " + f.getName() + " @ " + f.getEntryPoint() + " -----");
					if (res != null && res.decompileCompleted() && res.getDecompiledFunction() != null) {
						pw.println(res.getDecompiledFunction().getC());
					}
					else {
						pw.println("(decompile failed) " + (res == null ? "" : res.getErrorMessage()));
					}
					count++;
				}
			}
			finally {
				decomp.dispose();
			}
		}
		finally {
			pw.close();
		}
		println("Wrote " + out.getAbsolutePath());
	}
}