/***************************************************************************
    name                 : wmartag.c
    version              : 0.2
    description          : Display or modify a WMA file's title and author
                           tags.
    usage                : wmartag [-a "Author"] [-t "Title"] "File.wma"
                             -------------------
    begin                : 27 March, 2009
    copyright            : (C) 2009 by Kamran Riaz Khan
    email                : krkhan@inspirated.com
    website              : http://inspirated.com/
 **************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 **************************************************************************/

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define DEBUG
#define TFLAG 1
#define AFLAG 1 << 1
#define ASFCDOGUID \
	"75B22633-668E-11CF-A6D9-00AA0062CE6C"
/*
 * ASF Content Description Object GUID, refer ASF specification at:
 * http://www.microsoft.com/windows/windowsmedia/forpros/format/asfspec.aspx
 */

typedef unsigned char BYTE[1];
typedef unsigned char WCHAR[2];
typedef unsigned char WORD[2];
typedef unsigned char DWORD[4];
typedef unsigned char QWORD[8];
typedef unsigned char GUID[16];

typedef struct {
	/* guid is the _formatted_ 36 char string identifying ASF objects */
	char guid[37];
	/* rawguid is the 16 bytes of guid untouched */
	GUID rawguid;
	long size;
} ASFObject;

typedef struct {
	char *title;
	char *author;
	char *copyright;
	char *desc;
	char *rating;
} ContentInfo;

/* Wrapper function for fread */
void freadw(void *ptr, size_t size, size_t nmemb, FILE * stream)
{
	size_t result = fread(ptr, size, nmemb, stream);
	if (result != nmemb) {
		fputs("Read error\n", stderr);
		exit(1);
	}
}

/* Wrapper function for frwite */
void fwritew(void *ptr, size_t size, size_t nmemb, FILE * stream)
{
	size_t result = fwrite(ptr, size, nmemb, stream);
	if (result != nmemb) {
		fputs("Write error\n", stderr);
		exit(1);
	}
}

/* Convert low-endian bytestring into a long */
long letolong(unsigned char *str, unsigned int size)
{
	long value;
	int i;

	value = 0;
	for(i = size - 1; i >= 0; i--) {
		value <<= 8;
		value |= str[i];
	}

	return value;
}

/* Convert a long into a low-endian bytestring */
void longtole(long value, unsigned char *str, unsigned int size)
{
	long tmp;
	int i;

	for (i = 0; i < size; i++) {
		tmp = value;
		tmp <<= (sizeof(long) - 1 - i) * 8;
		tmp >>= (sizeof(long) - 1 - i) * 8;
		tmp >>= i * 8;
		str[i] = tmp;
	}
}

/* Convert a long into a low-endian WORD (2 bytes) */
void longtoword(long value, WORD str)
{
	longtole(value, str, 2);
}

/* Convert a long into a low-endian DWORD (4 bytes) */
void longtodword(long value, DWORD str)
{
	longtole(value, str, 4);
}

/* Convert a long into a low-endian QWORD (8 bytes) */
void longtoqword(long value, QWORD str)
{
	longtole(value, str, 8);
}

/* Convert a low-endian WORD (2 bytes) into a long */
long wordtolong(WORD word)
{
	return letolong(word, 2);
}

/* Convert a low-endian DWORD (4 bytes) into a long */
long dwordtolong(DWORD dword)
{
	return letolong(dword, 4);
}

/* Convert a low-endian QWORD (8 bytes) into a long */
long qwordtolong(QWORD qword)
{
	return letolong(qword, 8);
}

/* Convert a raw guid into its 36 char formatted representation */
void formatguid(GUID guid, char str[37])
{
	snprintf(str,
		 37,
		 "%02X%02X%02X%02X-%02X%02X-%02X%02X-"
		 "%02X%02X-%02X%02X%02X%02X%02X%02X",
		 guid[3], guid[2], guid[1], guid[0], guid[5], guid[4],
		 guid[7], guid[6], guid[8], guid[9], guid[10], guid[11],
		 guid[12], guid[13], guid[14], guid[15]);
}

/*
 * Convert a UTF-16 string into its ASCII equivalent
 * The returned string *must* be freed
 */
char *utf16toascii(char *ustr, size_t len)
{
	char *str = (char *) malloc(len / 2);
	int i;

	for (i = 0; i < len / 2; i++) {
		str[i] = ustr[i * 2];
	}

	return str;
}

/* Convert an ASCII string into its UTF-16 equivalent
 * The returned string *must* be freed
 */
char *asciitoutf16(char *str)
{
	size_t len = strlen(str) * 2 + 2;
	char *ustr = (char *) malloc(len);
	int i;

	memset(ustr, 0, len);
	for (i = 0; i < len; i += 2) {
		ustr[i] = str[i / 2];
	}

	return ustr;
}

/*
 * Fill an ASFObject from file
 * The cursor position will be modified
 */
ASFObject readobject(FILE * file)
{
	GUID rawguid;
	QWORD qword;
	ASFObject obj;

	freadw(rawguid, 1, 16, file);
	memcpy(obj.rawguid, rawguid, 16);
	formatguid(rawguid, obj.guid);

	freadw(qword, 1, 8, file);
	obj.size = qwordtolong(qword);

	return obj;
}

/*
 * Read UTF-16 bytes of length len from file, convert to ASCII and return
 * The returning string _must_ be freed
 */
char *readutf16(FILE *file, long len)
{
	char *ustr, *str;

	if(len) {
		ustr = (char *) malloc(len);
		freadw(ustr, 1, len, file);
		str = utf16toascii(ustr, len);
		free(ustr);
	} else {
		str = (char *) malloc(1);
		str[0] = '\0';
	}

	return str;
}

#ifdef DEBUG
/*
 * Print contents of an ASF Content Description Object pointed to by header
 * For debugging purposes only
 */
void printcontent(unsigned char *header)
{
	int i, j;
	long titlelen, authorlen, copyrightlen, desclen, ratinglen;

	printf("GUID: ");
	for (i = 0; i < 16; i++) {
		printf("%08X ", header[i]);
	}
	printf("\n");

	printf("SIZE: ");
	for (i = 16; i < 24; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");

	printf("TLEN: ");
	for (i = 24; i < 26; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	titlelen = wordtolong(&header[24]);

	printf("ALEN: ");
	for (i = 26; i < 28; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	authorlen = wordtolong(&header[26]);

	printf("CLEN: ");
	for (i = 28; i < 30; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	copyrightlen = wordtolong(&header[28]);

	printf("DLEN: ");
	for (i = 30; i < 32; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	desclen = wordtolong(&header[30]);

	printf("RLEN: ");
	for (i = 32; i < 34; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	ratinglen = wordtolong(&header[32]);

	j = 34;
	printf("TTLE: ");
	for (i = j; i < j + titlelen; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	j += titlelen;

	printf("ATHR: ");
	for (i = j; i < j + authorlen; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	j += authorlen;

	printf("CRHT: ");
	for (i = j; i < j + copyrightlen; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	j += copyrightlen;

	printf("DESC: ");
	for (i = j; i < j + desclen; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
	j += desclen;

	printf("RTNG: ");
	for (i = j; i < j + ratinglen; i++) {
		printf("%02X ", header[i]);
	}
	printf("\n");
}
#endif

/*
 * Read values in the ASF Content Description Object pointed to by the
 * cursor position in file
 */
ContentInfo readcontent(FILE *file)
{
	ContentInfo info;
	WORD word;
	long titlelen, authorlen, copyrightlen, desclen, ratinglen;

	freadw(word, 1, 2, file);
	titlelen = wordtolong(word);

	freadw(word, 1, 2, file);
	authorlen = wordtolong(word);

	freadw(word, 1, 2, file);
	copyrightlen = wordtolong(word);

	freadw(word, 1, 2, file);
	desclen = wordtolong(word);

	freadw(word, 1, 2, file);
	ratinglen = wordtolong(word);

	info.title = readutf16(file, titlelen);
	info.author = readutf16(file, authorlen);
	info.copyright = readutf16(file, copyrightlen);
	info.desc = readutf16(file, desclen);
	info.rating = readutf16(file, ratinglen);

	return info;
}

/*
 * Replace an ASF Content Description Object in the file
 * The info and obj structures provide details about the old object
 * The file cursor *must* be  already at the end of old object that
 * needs to be replaced
 */
void replacecontent(ContentInfo info, ASFObject obj, FILE *file)
{
	ASFObject newobj;
	QWORD qword;
	WORD word;
	long offset, bufsize, hoffset, hsize,
	     titlelen, authorlen, copyrightlen, desclen, ratinglen;
	unsigned char *buffer, *header;
	char *ustr;

	titlelen = strlen(info.title) ? strlen(info.title) * 2 + 2 : 0;
	authorlen = strlen(info.author) ? strlen(info.author) * 2 + 2 : 0;
	copyrightlen =
		strlen(info.copyright) ? strlen(info.copyright) * 2 + 2 : 0;
	desclen = strlen(info.desc) ? strlen(info.desc) * 2 + 2 : 0;
	ratinglen = strlen(info.desc) ? strlen(info.rating) * 2 + 2 : 0;

	memcpy(newobj.guid, obj.guid, 37);
	memcpy(newobj.rawguid, obj.rawguid, 16);
	newobj.size
	    = 34
	    + titlelen
	    + authorlen
	    + copyrightlen
	    + desclen
	    + ratinglen;

	offset = ftell(file);
	fseek(file, 0, SEEK_END);
	bufsize = ftell(file) - offset;
	fseek(file, offset, SEEK_SET);

	buffer = (unsigned char *) malloc(bufsize);
	if (buffer == NULL) {
		fputs("Memory error", stderr);
		exit(2);
	}
	freadw(buffer, 1, bufsize, file);

	header = (unsigned char *) malloc(newobj.size);
	memset(header, 0, newobj.size);

	memcpy(header, obj.rawguid, 16);
	hoffset = 16;

	longtoqword(newobj.size, qword);
	memcpy(header + hoffset, qword, 8);
	hoffset += 8;

	longtoword(titlelen, word);
	memcpy(header + hoffset, word, 2);
	hoffset += 2;

	longtoword(authorlen, word);
	memcpy(header + hoffset, word, 2);
	hoffset += 2;

	longtoword(copyrightlen, word);
	memcpy(header + hoffset, word, 2);
	hoffset += 2;

	longtoword(desclen, word);
	memcpy(header + hoffset, word, 2);
	hoffset += 2;

	longtoword(ratinglen, word);
	memcpy(header + hoffset, word, 2);
	hoffset += 2;

	ustr = asciitoutf16(info.title);
	memcpy(header + hoffset, ustr, titlelen);
	free(ustr);
	hoffset += titlelen;

	ustr = asciitoutf16(info.author);
	memcpy(header + hoffset, ustr, authorlen);
	free(ustr);
	hoffset += authorlen;

	ustr = asciitoutf16(info.copyright);
	memcpy(header + hoffset, ustr, copyrightlen);
	free(ustr);
	hoffset += copyrightlen;

	ustr = asciitoutf16(info.desc);
	memcpy(header + hoffset, ustr, desclen);
	free(ustr);
	hoffset += desclen;

	ustr = asciitoutf16(info.rating);
	memcpy(header + hoffset, ustr, ratinglen);
	free(ustr);

#ifdef DEBUG
	/* New header */
	printcontent(header);
#endif

	fseek(file, offset - obj.size, SEEK_SET);
	fwritew(header, 1, newobj.size, file);

	/* Write the buffered contents after the header */
	fwritew(buffer, 1, bufsize, file);
	/* Truncate the file, ftruncate is not ANSI C but
	 * works better than using another file for content modification */
	ftruncate(fileno(file), ftell(file));

	/* Update header count value in the beginning */
	fseek(file, 16, SEEK_SET);
	freadw(qword, 1, 8, file);
	hsize = qwordtolong(qword);
	if(newobj.size < obj.size) {
		hsize -= obj.size - newobj.size;
	} else {
		hsize += newobj.size - obj.size;
	}
	longtoqword(hsize, qword);
	fseek(file, 16, SEEK_SET);
	fwritew(qword, 1, 8, file);

	/* Seek to the beginning of next ASF object for the loop in main */
	fseek(file, offset - obj.size + newobj.size, SEEK_SET);

	printf("Tags updated successfully\n");

	free(header);
	free(buffer);
}

/* Free dynamically allocated strings in a ContentInfo structure */
void freeinfo(ContentInfo info)
{
	free(info.title);
	free(info.author);
	free(info.copyright);
	free(info.desc);
	free(info.rating);
}

int main(int argc, char *argv[])
{
	FILE *file;
	ASFObject obj;
	DWORD dword;
	BYTE byte;
	char *title, *author;
	long hcount;
	int i, flags;

	if(argc < 2 || argc > 6) {
		fputs("Usage: wmartag"
			"[-a \"Author\"] [-t \"Title\"] \"File.wma\"\n",
			stderr);
		exit(1);
	}

	file = fopen(argv[argc - 1], "r+");
	if (file == NULL) {
		fputs("File error\n", stderr);
		exit(1);
	}

	flags = 0;
	for (i = 1; i < argc - 1; i++) {
		if (!strcmp("-t", argv[i])) {
			title = (char *) malloc(strlen(argv[i + 1]));
			strcpy(title, argv[i + 1]);
			flags |= TFLAG;
		}
		if (!strcmp("-a", argv[i])) {
			author = (char *) malloc(strlen(argv[i + 1]));
			strcpy(author, argv[i + 1]);
			flags |= AFLAG;
		}
	}

	/* Main head object */
	obj = readobject(file);

#ifdef DEBUG
	printf("GUID:\t%s\t[%ld]\n", obj.guid, obj.size);
#endif

	freadw(dword, 1, 4, file);
	hcount = dwordtolong(dword);

	if (hcount < 1) {
		fputs("Invalid header count\n", stderr);
		exit(1);
	}

	/* Read reserved bytes */
	freadw(byte, 1, 1, file);
	freadw(byte, 1, 1, file);

	/* Remaining header objects */
	for (i = 0; i < hcount; i++) {
		obj = readobject(file);
#ifdef DEBUG
		printf("GUID:\t%s\t[%ld]\n", obj.guid, obj.size);
#endif
		if (!strcmp(obj.guid, ASFCDOGUID)) {
			ContentInfo info = readcontent(file);

			if (flags & TFLAG) {
				free(info.title);
				info.title = (char *) malloc(strlen(title));
				strcpy(info.title, title);
			}
			if (flags & AFLAG) {
				free(info.author);
				info.author
					= (char *) malloc(strlen(author));
				strcpy(info.author, author);
			}

			if (flags) {
				replacecontent(info, obj, file);
			} else {
				printf("Title: %s\n", info.title);
				printf("Author: %s\n", info.author);
				printf("Copyright: %s\n", info.copyright);
				printf("Description: %s\n", info.desc);
				printf("Rating: %s\n", info.rating);
			}

			freeinfo(info);
		} else {
			fseek(file, obj.size - 24, SEEK_CUR);
		}
	}

	fclose(file);

	if (flags & TFLAG) {
		free(title);
	}
	if (flags & AFLAG) {
		free(author);
	}

	return 0;
}