/*
 * Excel 97-2010 VBA Password Remover
 * Written by Kévin Subileau - www.kevinsubileau.fr
 * 
 * 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.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <iostream>
#include <cerrno>
#include <fstream>
#ifdef __MINGW32__
#include <stdlib.h>
#endif

#ifndef NO_XLSM_SUPPORT
	extern "C" {
		#include <zip.h>
	}
	#define XLSM_VBA_FILE "xl/vbaProject.bin"
#endif

using namespace std;

bool fileExist(const char * filename)
{
	ifstream ifile(filename);
	return ifile;
}

bool fileExist(string & filename)
{
	return fileExist(filename.c_str());
}

void copyFile(const char * src, const char * dest, bool binaryMode = true)
{
	std::ifstream  srcstream(src,ios_base::in | (binaryMode?ios_base::binary:ios_base::out)); 
	if (!srcstream.good())
		throw "Error opening file for copy!";
	
	std::ofstream  dststream(dest, ios_base::out | (binaryMode?ios_base::binary:ios_base::out));

	dststream << srcstream.rdbuf();
}

void copyFile(string src, string dest)
{
	copyFile(src.c_str(), dest.c_str());
}

#ifndef NO_XLSM_SUPPORT

bool patchXLSM(const char* xlsmFile, const char * buffer, streamsize bufsize)
{

	int visu = 0;
	struct zip * f_zip=NULL;
	struct zip_source * n_zip=NULL;
	int err = 0;

	f_zip=zip_open(xlsmFile,ZIP_CREATE,&err);
	
	if(err != ZIP_ER_OK)
	{
		char buf_erreur[2048];
		zip_error_to_str(buf_erreur, sizeof buf_erreur, err, errno);
		printf("Error %d : %s\n",err, buf_erreur);
		return false;
	}
	
	if(f_zip==NULL)
	{
		cerr << "Error while opening file " << xlsmFile << endl;
		return false;
	}
	
	if((n_zip=zip_source_buffer(f_zip, buffer, bufsize, 1)) == NULL)
	{
		printf("%s\n", zip_strerror(f_zip));
		zip_close(f_zip);
		f_zip = NULL;
		return false;
	}

	visu=zip_name_locate(f_zip,XLSM_VBA_FILE,0);
	if (visu!=-1)
	{
		if(zip_replace(f_zip,visu,n_zip) == -1)
		{
		  printf("%s\n", zip_strerror(f_zip));
		  zip_close(f_zip);
		  f_zip = NULL;
		  zip_source_free(n_zip);
		  n_zip = NULL;
		  return false;
		}
	}
	else
	{
		return false;
	}

	zip_close(f_zip);
	f_zip = NULL;
	n_zip = NULL;

	return true;

}

int readXLSM(const char* xlsmFile, char ** contents) {
	int err = 0;
	 zip *z=NULL;

	z = zip_open(xlsmFile, ZIP_CHECKCONS, &err);

	if(err != ZIP_ER_OK)
	{
		char buf_erreur[2048];
		zip_error_to_str(buf_erreur, sizeof buf_erreur, err, errno);
		printf("Error %d : %s\n",err, buf_erreur);
		return -1;
	}

    //Search for the file of given name
    struct zip_stat st;
    zip_stat_init(&st);
    zip_stat(z, XLSM_VBA_FILE, 0, &st);

    //Alloc memory for its uncompressed contents
    *contents = new char[st.size];

    //Read the compressed file
    zip_file *f = zip_fopen(z, XLSM_VBA_FILE, 0);
    zip_fread(f, *contents, st.size);
    zip_fclose(f);

    //And close the archive
    zip_close(z);
    
    return st.size;
}
#endif

int searchBytes(const char * buffer, int bufsize, string needle) {
	
	string haystack(buffer, bufsize);
	size_t n = haystack.rfind(needle);
	if (n == string::npos)
		return -1;
	else
		return n;
}

int hackdata(char * const buffer, streamsize bufsize)
{
	/* Search and replace the DPB key */
	int pos = searchBytes(buffer, bufsize, string("DPB=\"", 3));
	if(pos == -1)
		return -1;
	buffer[pos+2] = 'X';
	return 0;
}

int excel97 (string filename)
{
	char * buffer;
	streamsize bufsize;
	ifstream filein(filename.c_str(), ios_base::binary|ios_base::in|ios::ate);
	
	/* Get file size */
	bufsize = filein.tellg();
	filein.seekg (0, ios::beg);
	
	/* Read to buffer */
	cout << "Reading file..." << endl;
	buffer = new char [bufsize];
	filein.read (buffer, bufsize);
	filein.close();
	
	/* HACK ! */
	cout << "Doing the magic..." << endl;
	if(hackdata(buffer,bufsize) != 0) {
		cerr << "Error : can't find the DPB key :(. ";
		if(searchBytes(buffer, bufsize, string("DPX", 3)) != -1)
			cerr << "It's seems you have already hacked this file ! ";
		cerr << endl;
		delete[] buffer;
		return 1;
	}
	
	/* Write result */
	cout << "Patching file..." << endl;
	ofstream fileout(filename.c_str(), ios_base::binary|ios_base::out);
	fileout.write (buffer, bufsize);
	fileout.close();
    
	cout << "Success !" << endl;
	cout << "Now open the file, go to VBA Editor (Alt+F11), click yes on the warning message, choose the menu command Tools->VBAProject Properties, navigate to the Protection tab, and change the password." << endl;
	delete[] buffer;
	return 0;
}

#ifndef NO_XLSM_SUPPORT
int excel2010 (string filename)
{
	char * buffer;
	cout << "Extracting data..." << endl;
	streamsize bufsize = readXLSM(filename.c_str(), &buffer);
	
	if(bufsize == -1) {
		cerr << "Error while reading workbook content. " << endl;
		return 1;
	}
	
	cout << "Doing the magic..." << endl;
	if(hackdata(buffer,bufsize) != 0) {
		cerr << "Error : can't find the DPB key :(. ";
		if(searchBytes(buffer, bufsize, string("DPX", 3)) != -1)
			cerr << "It's seems you have already hacked this file ! ";
		cerr << endl;
		return 2;
	}
    
	cout << "Patching file..." << endl;
	if(!patchXLSM(filename.c_str(), buffer, bufsize)) {
		cerr << "Error while patching workbook content. " << endl;
		return 3;
	}
	
	cout << "Success !" << endl;
	cout << "Now open the file, go to VBA Editor (Alt+F11), click yes on the warning message, choose the menu command Tools->VBAProject Properties, navigate to the Protection tab, and change the password." << endl;
    return 0;
}
#endif

int main(int argc, char **argv)
{
	if (argc < 2)
	{
		puts("You must specify the file you want to unlock !");
    #ifdef __MINGW32__
		system("pause");
	#endif
		return 1;
	}
	string targetFile = argv[1];
	
	if(!fileExist(targetFile)) {
		cerr << "Error : File not found" << endl;
    #ifdef __MINGW32__
		system("pause");
	#endif
		return 3;
	}
	
	int extdotpos = targetFile.find_last_of(".");
	string ext = targetFile.substr(extdotpos + 1);
	int ret;
	if(ext == "xls") {
		cout << "Working on an Excel 97 - 2003 file" << endl;
		cout << "Backing up file..." << endl;
		copyFile(targetFile, targetFile.substr(0, extdotpos) + "_bak.xls");
		
		ret = excel97(targetFile);
	} 
#ifndef NO_XLSM_SUPPORT
	else if(ext == "xlsm") {
		cout << "Working on an Excel 2010 file" << endl;
		cout << "Backing up file..." << endl;
		copyFile(targetFile, targetFile.substr(0, extdotpos) + "_bak.xlsm");
		
		ret = excel2010(targetFile);
	}
	else if(ext == "xlsx") {
		cerr << "Error : This Excel 2010 file format does not support macro. Please specify a xlsm file instead." << endl;
		ret = 4;
	}
#else
	else if(ext == "xlsm" || ext == "xlsx") {
		cerr << "Error : Excel 2010 workbooks are not supported." << endl;
		ret = 10;
	}
#endif
	else {
		cerr << "Error : Unrecognized file extension" << endl;
		ret = 4;
	}
    #ifdef __MINGW32__
		system("pause");
	#endif
	return ret;
}