/*
 *  FreeMe main.c -- mostly a wma/asf file processor, with DRM part
 *  put off into msdrm.c
 */

#include <windows.h>
#include <sys/types.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include "msdrm.h"

typedef unsigned int uint32_t;
typedef unsigned __int64 uint64_t;

#pragma pack(1)


typedef struct chunk_save_st {
	GUID guid;
	uint64_t len;
	void *data;
	struct chunk_save_st *next;
} CHUNKSAVE;

typedef int (*GUIDHANDLER) (FILE * fp, CHUNKSAVE * savep);

typedef struct guidaction_st {
	GUID *guid;
	char *name;
	GUIDHANDLER fn;
} GUIDACTION;

typedef struct fileheader_st {
	GUID clientGUID;
	uint64_t filesize;
	uint64_t fileCreateTime;
	uint64_t numPackets;
	uint64_t timeAtEnd;
	uint64_t playDuration;
	uint32_t timeAtStart;
	uint32_t unknown1;
	uint32_t unknown2;
	uint32_t packetSize;
	uint32_t packetSize2;
	uint32_t uncompressedSize;
} FILEHEADER;

int handle_chunk(FILE * fp, CHUNKSAVE * chunk);
int handle_header(FILE * fp, CHUNKSAVE * chunk);
int handle_file_header(FILE * fp, CHUNKSAVE * chunk);
int handle_data(FILE * fp, CHUNKSAVE * chunk);
int handle_stream_header(FILE * fp, CHUNKSAVE * chunk);
int handle_copy(FILE * fp, CHUNKSAVE * chunk);
int handle_drmv1(FILE * fp, CHUNKSAVE * chunk);
int handle_drmv2(FILE * fp, CHUNKSAVE * chunk);

struct globalinfo_st globalinfo;

GUID HeaderGUID = { 0x75b22630, 0x668e, 0x11cf, {0xa6, 0xd9,
						 0x00, 0xaa, 0x00, 0x62,
						 0xce, 0x6c}
};

GUID DataGUID = { 0x75b22636, 0x668e, 0x11cf, {0xa6, 0xd9,
					       0x00, 0xaa, 0x00, 0x62,
					       0xce, 0x6c}
};

GUID FileHeaderGUID = { 0x8cabdca1, 0xa947, 0x11cf, {0x8e, 0xe4,
						     0x00, 0xc0, 0x0c,
						     0x20, 0x53, 0x65}
};

GUID StreamHeaderGUID = { 0xb7dc0791, 0xa9b7, 0x11cf, {0x8e, 0xe6,
						       0x00, 0xc0, 0x0c,
						       0x20, 0x53, 0x65}
};

GUID AudioStreamGUID = { 0xf8699e40, 0x5b4d, 0x11cf, {0xa8, 0xfd,
						      0x00, 0x80, 0x5f,
						      0x5c, 0x44, 0x2b}
};

GUID Unknown1GUID = { 0x5fbf03b5, 0xa92e, 0x11cf, {0x8e, 0xe3,
						   0x00, 0xc0, 0x0c, 0x20,
						   0x53, 0x65}
};

GUID Unknown2GUID = { 0x86d15240, 0x311d, 0x11d0, {0xa3, 0xa4,
						   0x00, 0xa0, 0xc9, 0x03,
						   0x48, 0xf6}
};

GUID DRMv2HeaderGUID = { 0x298ae614, 0x2622, 0x4c17, {0xb9, 0x35,
						      0xda, 0xe0, 0x7e,
						      0xe9, 0x28, 0x9c}
};

GUID DRMv1HeaderGUID = { 0x2211b3fb, 0xbd23, 0x11d2, {0xb4, 0xb7,
						      0x00, 0xa0, 0xc9,
						      0x55, 0xfc, 0x6e}
};

GUID ContentDescrGUID = { 0x75b22633, 0x668e, 0x11cf, {0xa6, 0xd9,
						       0x00, 0xaa, 0x00,
						       0x62, 0xce, 0x6c}
};

GUID PropertyListGUID = { 0xd2d0a440, 0xe307, 0x11d2, {0x97, 0xf0,
						       0x00, 0xa0, 0xc9,
						       0x5e, 0xa8, 0x50}
};

GUIDACTION known_guids[] = {
	{&HeaderGUID, "Header", handle_header},
	{&DataGUID, "Data", handle_data},
	{&FileHeaderGUID, "File Header", handle_file_header},
	{&StreamHeaderGUID, "Stream Header", handle_copy},
	{&Unknown1GUID, "Header subchunk - unknown 1", handle_copy},
	{&Unknown2GUID, "Header subchunk - unknown 2", handle_copy},
	{&DRMv2HeaderGUID, "DRMV2 ContentHeader", handle_drmv2},
	{&DRMv1HeaderGUID, "DRMv1 header", handle_drmv1},
	{&ContentDescrGUID, "Content Description", handle_copy},
	{&PropertyListGUID, "Property List", handle_copy},
	{NULL, NULL, NULL}
};


void error_exit(char *msg)
{
	fprintf(stderr, "%s\n", msg);
	fprintf(stderr, "\n   Press <ENTER> to acknowledge error.\n");
	getchar();
	exit(1);
}

void printwcs(wchar_t * msg)
{
	int len, rval;
	char *buff;

	len = wcslen(msg);
	if ((buff = malloc(len + 1)) == NULL)
		error_exit("Memory allocation failed in printwcs");

	rval = WideCharToMultiByte(CP_ACP, 0, msg, len + 1, buff, len + 1,
				   NULL, NULL);
	if (rval == 0)
		error_exit("WideCharToMultiByte failed in printwcs");

	fputs(buff, stderr);
	free(buff);
}


GUIDHANDLER find_guid(GUID * guid)
{
	GUIDHANDLER handler = NULL;
	GUIDACTION *curr;

	curr = known_guids;
	while (curr->guid != NULL) {
		if (memcmp(curr->guid, guid, sizeof(GUID)) == 0) {
			handler = curr->fn;
			break;
		}
		curr++;
	}

	return handler;
}


int handle_copy(FILE * fp, CHUNKSAVE * chunk)
{
	unsigned long datalen = (unsigned long) (chunk->len - 24);

	chunk->data = malloc(datalen);

	if (chunk->data == NULL)
		error_exit("Memory allocation failed in handle_copy");

	if (fread(chunk->data, datalen, 1, fp) != 1) {
		free(chunk->data);
		chunk->data = NULL;
		return 0;
	}

	return 1;
}


int handle_drmv1(FILE * fp, CHUNKSAVE * chunk)
{
	unsigned long toskip = (unsigned long) (chunk->len - 24);

	fseek(fp, toskip, SEEK_CUR);
	globalinfo.hasV1header = 1;
	if (globalinfo.verbose)
		fprintf(stderr, "Found DRMv1 header object.\n");
	return 1;
}


int handle_drmv2(FILE * fp, CHUNKSAVE * chunk)
{
	unsigned long datalen = (unsigned long) (chunk->len - 24);
	unsigned short *data = malloc(datalen);

	if (globalinfo.verbose)
		fprintf(stderr, "Found DRMv2 header object.\n");

	if (data == NULL)
		error_exit("Memory allocation in handle_drmv2 failed.");

	fseek(fp, 6, SEEK_CUR);
	if (fread(data, datalen - 6, 1, fp) != 1)
		error_exit("Data read in handle_drmv2 failed.");

	data[(datalen - 6) / 2] = L'\0';

	globalinfo.kid = get_element(L"KID", data);
	globalinfo.hasV2header = 1;

	if (globalinfo.verbose) {
		if (globalinfo.kid == NULL) {
			fprintf(stderr,
				"KID not found in header object!\n");
		} else {
			fprintf(stderr, "Found KID (");
			printwcs(globalinfo.kid);
			fprintf(stderr, ")\n");
		}
	}

	free(data);
	return 1;
}


int handle_packet(FILE * fp, int packetlen)
{
	struct packethead_st {
		uchar id;
		short unknown1;
		uchar flags;
		uchar segTypeID;
	} *info;
	uchar *data;
	int flagoffset = 13;
	int objlen;
	int dataoffset;
	int rval = 0;

	if ((data = malloc(packetlen)) == NULL)
		error_exit("Memory allocation failed in handle_packet.");

	if (fread(data, packetlen, 1, fp) != 1)
		goto exit;

	info = (struct packethead_st *) data;
	if (info->id != 0x82)
		error_exit("Unknown packet id - don't know what to do!");

	if (info->flags & 0x40)
		flagoffset += 2;

	if (info->flags & 0x10)
		flagoffset += 2;
	else if (info->flags & 0x08)
		flagoffset += 1;

	if (info->flags & 0x01)
		flagoffset += 1;

	if (info->segTypeID == 0x55)
		flagoffset += 1;
	else if (info->segTypeID == 0x59)
		flagoffset += 2;
	else if (info->segTypeID == 0x5d)
		flagoffset += 4;

	if (data[flagoffset] != 8)
		error_exit
		    ("Flag says grouping - don't know how to do this!");

	dataoffset = flagoffset + 9;
	if (info->flags & 0x01)
		error_exit("Need the data_length field - don't know how!");

	objlen = *((int *) &data[flagoffset + 1]);

	MSDRM_decr_packet(data + dataoffset, objlen,
			  globalinfo.content_key);

	fwrite(data, packetlen, 1, stdout);
	rval = 1;

exit:
	free(data);
	return rval;
}


int handle_data(FILE * fp, CHUNKSAVE * chunk)
{
	struct datahead_st {
		GUID unknownGUID;
		uint64_t numPackets;
		uchar unknown[2];
	} datahead;
	int packetcount = 0;
	int lastpercent = -1;

	if (fread(&datahead, sizeof(datahead), 1, fp) != 1)
		return 0;

	fwrite(chunk, 24, 1, stdout);
	fwrite(&datahead, sizeof(datahead), 1, stdout);

	if (globalinfo.verbose) {
		fprintf(stderr, "Starting to process data packets\n");
		fprintf(stderr, "%d packets of length %d\n",
			globalinfo.numpackets, globalinfo.packetlen);
	}

	while (handle_packet(fp, globalinfo.packetlen)) {
		packetcount++;
		if (globalinfo.numpackets != 0) {
			int percent, i;
			percent =
			    ((packetcount * 200) / globalinfo.numpackets +
			     1) / 2;
			if (percent != lastpercent) {
				fprintf(stderr, "|");
				for (i = 0; i < percent / 2; i++)
					fprintf(stderr, "#");
				for (; i < 50; i++)
					fprintf(stderr, " ");
				fprintf(stderr, "|  ");
				fprintf(stderr, "%3d%%\r", percent);
				lastpercent = percent;
			}
		}
	}
	fprintf(stderr, "\n");

	return 0;
}


int handle_header(FILE * fp, CHUNKSAVE * chunk)
{
	struct header_st {
		int numchunks;
		short unknown;
	} header;
	int i;
	CHUNKSAVE *subchunk = NULL;
	int savecount = 0;
	CHUNKSAVE *head = NULL, *tail = NULL;
	FILEHEADER *fileheader;
	uint64_t bytesremoved = 0;

	if (fread(&header, sizeof(header), 1, fp) != 1)
		return 0;

	for (i = 0; i < header.numchunks; i++) {
		if (subchunk == NULL) {
			if ((subchunk = malloc(sizeof(CHUNKSAVE))) == NULL)
				error_exit
				    ("Memory allocation failed in handle_header");
		}
		if (!handle_chunk(fp, subchunk))
			return 0;
		if (subchunk->data != NULL) {
			if (tail == NULL)
				head = subchunk;
			else
				tail->next = subchunk;
			subchunk->next = NULL;
			tail = subchunk;

			if ((subchunk = malloc(sizeof(CHUNKSAVE))) == NULL)
				error_exit
				    ("Memory allocation failed in handle_header");

			savecount++;
		} else {
			bytesremoved += subchunk->len;
		}
	}

	if (globalinfo.fileheader == NULL) {
		error_exit("Didn't see file header!");
	} else {
		CHUNKSAVE *currchunk, *nextchunk;

		if (!globalinfo.hasV2header) {
			if (globalinfo.hasV1header)
				error_exit
				    ("This file is version 1 protected, not version 2.");
			else
				error_exit
				    ("This file doesn't seem to be protected!");
		}

		if (!globalinfo.kid)
			error_exit
			    ("Version 2 protected, but no KID found!");

		if (globalinfo.verbose)
			fprintf(stderr, "Starting to look for license.\n");

		globalinfo.content_key = MSDRM_init(globalinfo.kid);
		if (globalinfo.content_key == NULL)
			error_exit("Couldn't find a valid license!");

		if (freopen(globalinfo.ofname, "wb", stdout) == NULL)
			error_exit("Couldn't open output file.");

		if (globalinfo.verbose)
			fprintf(stderr, "Opened output file <%s>\n",
				globalinfo.ofname);

		currchunk = head;
		fileheader =
		    (FILEHEADER *) ((CHUNKSAVE *) globalinfo.fileheader)->
		    data;
		fileheader->filesize -= bytesremoved;
		globalinfo.packetlen = fileheader->packetSize;
		globalinfo.numpackets = (int) fileheader->numPackets;
		header.numchunks = savecount;
		chunk->len -= bytesremoved;
		fwrite(chunk, 24, 1, stdout);
		fwrite(&header, sizeof(header), 1, stdout);
		while (currchunk != NULL) {
			fwrite(currchunk, 24, 1, stdout);
			fwrite(currchunk->data,
			       (unsigned long) (currchunk->len - 24), 1,
			       stdout);
			nextchunk = currchunk->next;
			free(currchunk->data);
			free(currchunk);
			currchunk = nextchunk;
		}
	}

	return 1;
}


int handle_file_header(FILE * fp, CHUNKSAVE * chunk)
{
	unsigned long datalen = (unsigned long) (chunk->len - 24);

	if ((chunk->data = malloc(datalen)) == NULL)
		error_exit
		    ("Data allocation failed in handle_file_header.");

	if (fread(chunk->data, datalen, 1, fp) != 1) {
		free(chunk->data);
		chunk->data = NULL;
		return 0;
	}

	globalinfo.fileheader = chunk;

	return 1;
}


int handle_chunk(FILE * fp, CHUNKSAVE * chunk)
{
	GUIDHANDLER handler = NULL;
	int retval = 0;

	chunk->data = NULL;

	if (fread(&chunk->guid, sizeof(chunk->guid), 1, fp) != 1)
		return 0;

	if (fread(&chunk->len, sizeof(chunk->len), 1, fp) != 1)
		return 0;

	handler = find_guid(&chunk->guid);

	if (handler != NULL) {
		retval = (handler) (fp, chunk);
	} else {
		long toskip = (long) (chunk->len - 24);
		retval = (fseek(fp, toskip, SEEK_CUR) == 0);
		retval = 1;
	}

	return retval;
}


int main(int argc, char *argv[])
{
	CHUNKSAVE chunk;
	int more = 1;
	FILE *ifp;
	char *fnamestart;
	static char ofname[1000];

	globalinfo.verbose = 0;
	globalinfo.fileheader = NULL;
	globalinfo.kid = NULL;
	globalinfo.hasV1header = 0;
	globalinfo.hasV2header = 0;
	globalinfo.ofname = ofname;

	if ((argc < 2) || (argc > 3))
		error_exit("Usage: FreeMe [-v] protectedfile");

	if ((strcmp(argv[1], "-v") != 0) && (argc == 3))
		error_exit("Usage: FreeMe [-v] protectedfile");

	if (argc == 3)
		globalinfo.verbose = 1;

	if ((ifp = fopen(argv[argc - 1], "rb")) == NULL) {
		sprintf(ofname, "Couldn't open input file (%s)",
			argv[argc - 1]);
		error_exit(ofname);
	}

	ofname[0] = '\0';
	if ((fnamestart = strrchr(argv[argc - 1], '\\')) != NULL) {
		memcpy(ofname, argv[argc - 1],
		       fnamestart - argv[argc - 1] + 1);
		ofname[fnamestart - argv[argc - 1] + 1] = '\0';
	}
	strcat(ofname, "Freed-");
	strcat(ofname, (fnamestart ? fnamestart + 1 : argv[argc - 1]));

	while (more)
		more = handle_chunk(ifp, &chunk);

	return 0;
}