ExtractDATEmbeds.cpp

//--------------------------------------------//
// ExtractDATEmbeds                           //
// License: Public Domain (www.unlicense.org) //
//--------------------------------------------//

#include <stdlib.h>
#include <stdio.h>
#include <stdarg.h>
#include <string.h>
#include <string>
#ifdef _WIN32
#include <windows.h>
#endif

typedef unsigned char Bit8u;
#ifndef _MSC_VER
typedef unsigned long long Bit64u;
#else
typedef unsigned __int64 Bit64u;
#ifndef va_copy
#define va_copy(d,s) ((d) = (s))
#endif
#endif

static void LogErr(const char* fmt, ...) { va_list ap; va_start(ap, fmt); vfprintf(stderr, fmt, ap); va_end(ap); fflush(stderr); }
static void Log(const char* fmt, ...) { va_list ap; va_start(ap, fmt); vfprintf(stdout, fmt, ap); va_end(ap); fflush(stdout); }

#ifdef NDEBUG
#define ZIP_ASSERT(cond)
#else
#define ZIP_ASSERT(cond) (void)((cond) ? ((int)0) : *(volatile int*)0 |= 0xbad|fprintf(stderr, "FAILED ASSERT (%s)\n", #cond))
#endif

static FILE* fopen_utf8(const char* path, const char* mode)
{
	#ifdef _WIN32
	bool needw = false;
	for (const char* p = path; *p; p++) { if ((Bit8u)*p > 0x7F) { needw = true; break; } }
	if (needw)
	{
		WCHAR wpath[MAX_PATH+5], wmode[20], *pwmode = wmode;
		MultiByteToWideChar(CP_UTF8, 0, path, -1, wpath, MAX_PATH);
		for (const char* p = mode, *pEnd = p + 19; *p && p != pEnd; p++) *(pwmode++) = *p;
		*pwmode = '\0';
		return _wfopen(wpath, wmode);
	}
	#endif
	return fopen(path, mode);
}

enum EXml { XML_END = 0, XML_ELEM_START, XML_ELEM_END, XML_ELEM_SOLO, XML_TEXT };
static EXml XMLParse(char*& pStart, char*& pEnd)
{
	char* p = pStart;
	while (*p <= ' ') { if (!*p) return XML_END; p++; }
	pStart = p;
	if (*p == '<')
	{
		while (*p != '>') if (!*(p++)) return XML_END;
		pEnd = p + 1;
		if (pStart[1] == '?' || pStart[1] == '!') // skip XML instructions and syntax elements
		{
			// Skip past XML comments (`<!--  ... -->`)
			if (pStart[1] == '!' && pStart[2] == '-' && pStart[3] == '-') { p = pStart + 3; while (p[0] != '-' || p[1] != '-' || p[2] != '>') if (!*(p++)) return XML_END; pEnd = p + 3; }
			// Skip past DTD DOCTYPE (tags inside tags, it looks like `<!DOCTYPE name [ <!ELEMENT...>... ]>`
			else if (pStart[1] == '!') for (p = pStart; p != pEnd; p++) if (*p == '[') { while (*p) if (*(p++) == ']') { while (*p != '>') if (!*(p++)) return XML_END; return XMLParse(pStart = ++p, pEnd); } return XML_END; }
			return XMLParse(pStart = pEnd, pEnd);
		}
		return (p[-1] == '/' ? XML_ELEM_SOLO : (pStart[1] == '/' ? XML_ELEM_END : XML_ELEM_START));
	}
	while (*p != '<' && *p) p++;
	while (p[-1] <= ' ') p--;
	pEnd = p;
	return XML_TEXT;
}

static bool XMLMatchTag(char* pStart, char* pEnd, const char* tag, int tagLen, ...)
{
	ZIP_ASSERT(pStart[0] == '<' && pEnd[-1] == '>');
	int elemLen = (int)((pEnd -= (pEnd[-2] == '/' ? 2 : 1)) - (++pStart));
	if (elemLen < tagLen || memcmp(pStart, tag, tagLen) || (pStart[tagLen] != '>' && pStart[tagLen] != '/' && pStart[tagLen] != ' ')) return false;

	va_list apStart, ap;
	va_start(apStart, tagLen);
	va_copy(ap, apStart);
	int matchCount = 0, matchLens[64], n;
	for (const char *matchParam, **matchOut, **matchEndOut; (matchParam = va_arg(ap, const char*)) != NULL; matchCount++)
	{
		matchLens[matchCount] = (int)strlen(matchParam);
		*(matchOut = va_arg(ap, const char**)) = NULL;
		if ((matchEndOut = va_arg(ap, const char**)) != NULL) *matchEndOut = NULL;
	}
	va_end(ap);
	ZIP_ASSERT(matchCount < 64);

	Bit64u matchedBits = 0, bit;
	for (char* p = pStart + tagLen;;)
	{
		while (*p <= ' ') p++;
		if (p == pEnd) { va_end(apStart); return true; }
		const char *paramName = p;
		while (*p > ' ' && *p != '=') p++;
		const char *paramNameEnd = p, *val = NULL, *valEnd = NULL;
		while (*p <= ' ') p++;
		if (*p == '=')
		{
			p++;
			while (*p <= ' ') p++;
			if (*p == '"') { val = ++p; while (*p != '"' && p != pEnd) p++; valEnd = p; if (*p == '"') p++; }
			else { val = p; while (*p > ' ' && p != pEnd) p++; valEnd = p; }
		}

		va_copy(ap, apStart);
		for (n = 0, bit = 1; n != matchCount; n++, bit <<= 1)
		{
			const char *matchParam = va_arg(ap, const char*), **matchOut = va_arg(ap, const char**), **matchEndOut = va_arg(ap, const char**);
			if ((matchedBits & bit) || *matchParam != *paramName || matchLens[n] != (int)(paramNameEnd - paramName) || memcmp(paramName, matchParam, matchLens[n])) continue;
			*matchOut = val;
			if (matchEndOut) *matchEndOut = valEnd;
			matchedBits |= bit;
		}
		va_end(ap);
	}
}

static char* XMLLevel(char* p, EXml x = XML_END, char** textStart = NULL, char** textEnd = NULL)
{
	if (x && x != XML_ELEM_START) return p;
	int depth = (x ? 1 : 0), hadDepth = depth;
	for (char* pEnd = NULL;; p = pEnd)
	{
		if ((x = XMLParse(p, pEnd)) == XML_END) return NULL;
		if (x == XML_ELEM_START) { hadDepth = 1; ++depth; }
		if (x == XML_ELEM_END && --depth <= 0) return ((hadDepth && !depth) ? pEnd : NULL);
		if (x == XML_ELEM_SOLO && !depth) return pEnd;
		if (x == XML_TEXT && depth == 1 && textStart) { *textStart = p; *textEnd = pEnd; }
	}
}

int main(int argc, char *argv[])
{
	const char *xmlPath = NULL, *outPath = NULL;
	for (int i = 1; i < argc; i++)
	{
		if (!xmlPath && (argv[i][0] != '-' || !argv[i][1])) { xmlPath = argv[i]; continue; }
		if (!outPath && argv[i][0] != '-') { outPath = argv[i]; continue; }
		if (argv[i][0] != '-' || !argv[i][1] || argv[i][2]) goto argerr;
		switch (argv[i][1])
		{
			case 'i': if (xmlPath || ++i == argc) goto argerr; xmlPath = argv[i]; continue;
			case 'o': if (outPath || ++i == argc) goto argerr; outPath = argv[i]; continue;
		}
		argerr: LogErr("Unknown command line option '%s'.\n\n", argv[i]); goto help;
	}
	if (!xmlPath)
	{
		help:
		LogErr("%s v%s - Command line options:\n"
			"  [-i] <PATH>: Path to input XML file (or - to pass XML via stdin) (required)\n"
			"  [-o] <PATH>: Path to output file directory (defaults to XML file or current directory)\n"
			"\n", "ExtractDATEmbeds", "1.1");
		return 1;
	}

	bool useStdin = (!xmlPath || (xmlPath[0] == '-' && xmlPath[1] == '\0'));
	std::string xml;
	Log("Loading DAT XML from %s ...\n", (useStdin ? "<console input> (press CTRL+C to abort input)" : xmlPath));
	FILE* fXML = (useStdin ? stdin : fopen(xmlPath, "rb"));
	if (!fXML) { LogErr("Failed to open DAT XML %s\n\n", xmlPath); return 1; }
	if (!useStdin) { fseek(fXML, 0, SEEK_END); xml.reserve((size_t)ftell(fXML) + 1); fseek(fXML, 0, SEEK_SET); }

	char xmlreadbuf[2048];
	size_t readStep = (useStdin ? 1 : sizeof(xmlreadbuf));
	for (int rootTagOfs = 0, rootTagLen = 0, got = 0, ofs = 0;;)
	{
		if (ofs == got)
		{
			got = (int)fread(xmlreadbuf, 1, readStep, fXML);
			ofs = 0;
			if (got > 0) { if (useStdin) xml.reserve((xml.length()+65535)/32768*32768); xml.append(xmlreadbuf, got); }
		}
		int c = (ofs >= got ? 0 : xmlreadbuf[ofs++]);
		if (c == 0 || c == 4 || c == 23) // 4 is CTRL+D, 23 is CTRL+W
		{
			incompleteXml:
			LogErr("Received incomplete XML\n");
			return 1;
		}
		if (c != '>') continue;
		char *pXml = &xml[0], *pTag, *pTagEn;
		if (!rootTagLen)
		{
			EXml xmlRoot = XMLParse((pTag = pXml), pTagEn);
			if (xmlRoot == XML_END || pTagEn > pXml + xml.size() - (got - ofs)) continue;;
			if (xmlRoot == XML_ELEM_END || xmlRoot == XML_TEXT) goto incompleteXml;
			if (xmlRoot == XML_ELEM_SOLO) break;
			rootTagOfs = (int)(pTag + 1 - &xml[0]);
			while (pTag[rootTagLen+1] != '>' && pTag[rootTagLen+1] != ' ') rootTagLen++;
		}
		else if (!memcmp(pXml + xml.size() - 1 - (got - ofs) - rootTagLen, pXml + rootTagOfs, rootTagLen) && XMLLevel(pXml + rootTagOfs))
			{ xml.resize(xml.size() - (got - ofs) + 1); xml.back() = '\0'; break; } // force null terminate
	}
	fclose(fXML);

	const char *xmlPathEnd = xmlPath + (useStdin ? 0 : strlen(xmlPath));
	while (xmlPathEnd != xmlPath && *xmlPathEnd != '/' && *xmlPathEnd != '\\') xmlPathEnd--;

	std::string outBase, workPath, romOutPath;
	((outPath || xmlPathEnd == xmlPath) ? outBase.assign(outPath ? outPath : ".") : outBase.assign(xmlPath, (xmlPathEnd - xmlPath))) += '/';

	EXml x;
	for (char* pGame = &xml[0], *pGameEn = NULL; (x = XMLParse(pGame, pGameEn)) != XML_END; pGame = pGameEn)
	{
		char *gameName, *gameNameX;
		if (x == XML_TEXT || x == XML_ELEM_END || !XMLMatchTag(pGame, pGameEn, (pGame[1] == 'g' ? "game" : "machine"), (pGame[1] == 'g' ? 4 : 7), "name", &gameName, &gameNameX, NULL) || x == XML_ELEM_SOLO) continue;

		for (char* p = pGameEn, *pEnd, *pNext; p && (x = XMLParse(p, pEnd)) != XML_END && x != XML_ELEM_END && (pNext = XMLLevel(pEnd, x)) != NULL; p = pNext)
		{
			char *romName, *romNameX, *romData, *romDataX;
			if (!XMLMatchTag(p, pEnd, "rom", 3, "name", &romName, &romNameX, "data", &romData, &romDataX, NULL) || !romData) continue;

			romOutPath.assign(outBase).append(gameName, gameNameX - gameName).append(1, '.').append(romName, romNameX - romName);

			// Unescape XML characters (like &) and replace path separators in the ROM name with a dot character (because we don't want to extract into subdirectories)
			char* pRomOutPath = &romOutPath[0];
			for (char *p = pRomOutPath + outBase.size(), *pIn = p;;)
				if (!*pIn) { *p = '\0'; break; }
				else if (*pIn == '\\' || *pIn == '/') { *(p++) = '.'; pIn++; continue; }
				else if ((*(p++) = *(pIn++)) != '&') continue;
				else if (!strncmp(p, "amp;", 4)) pIn += 4;
				else if (!strncmp(p, "lt;", 3)) { p[-1] = '{'; pIn += 3; }
				else if (!strncmp(p, "gt;", 3)) { p[-1] = '}'; pIn += 3; }

			FILE* f = fopen_utf8(pRomOutPath, "wb");
			if (!f) { LogErr("  Unable to open output file: %s - Aborting\n", pRomOutPath); return 1; }
			Log("  Extracting: %s\n", pRomOutPath + outBase.length());

			static const Bit8u base64dec[256] = {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,62,0,0,0,63,52,53,54,55,56,57,58,59,60,61,0,0,0,0,0,0,0,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,0,0,0,0,0,0,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,0};
			for (Bit8u buf[3], lastlen = 0, *pData = (Bit8u*)romData, *pDataEnd = (Bit8u*)romDataX; pData != pDataEnd; pData += 4)
			{
				buf[0] = (base64dec[pData[0]] << 2);
				if ((pDataEnd - pData) != 1 && pData[1] != '=')
				{
					buf[0] |= ((base64dec[pData[1]] >> 4) & 0x3);
					buf[1] = (base64dec[pData[1]] << 4);
					if ((pDataEnd - pData) != 2 && pData[2] != '=')
					{
						buf[1] |= (base64dec[pData[2]] >> 2) & 0xf;
						buf[2] = (base64dec[pData[2]] << 6);
						if ((pDataEnd - pData) != 3 && pData[3] != '=') 
						{
							buf[2] |= (base64dec[pData[3]]);
							fwrite(buf, 3, 1, f);
							continue;
						}
						lastlen++;
					}
					lastlen++;
				}
				lastlen++;
				fwrite(buf, lastlen, 1, f);
				break;
			}

			fclose(f);
		}
	}
	Log("Finished.\n");

	return 0;
}